Prompt:
 + 1)([0, 1, 2]),
        [1, 2, 3],
    )
    self.assertEqual(
        (1 + rasp.tokens)([0, 1, 2]),
        [1, 2, 3],
    )

  def test_dunders_with_sop(self):
    self.assertEqual(
        (rasp.tokens + rasp.indices)([0, 1, 2]),
        [0, 2, 4],
    )
    self.assertEqual(
        (rasp.length - 1 - rasp.indices)([0, 1, 2]),
        [2, 1, 0],
    )
    self.assertEqual(
        (rasp.length * rasp.length)([0, 1, 2]),
        [9, 9, 9],
    )

  def test_logical_dunders(self):
    self.assertEqual(
        (rasp.tokens & True)([True, False]),
        [True, False],
    )
    self.assertEqual(
        (rasp.tokens & False)([True, False]),
        [False, False],
    )
    self.assertEqual(
        (rasp.tokens | True)([True, False]),
        [True, True],
    )
    self.assertEqual(
        (rasp.tokens | False)([True, False]),
        [True, False],
    )
    self.assertEqual(
        (True & rasp.tokens)([True, False]),
        [True, False],
    )
    self.assertEqual(
        (False & rasp.tokens)([True, False]),
        [False, False],
    )
    self.assertEqual(
        (True | rasp.tokens)([True, False]),
        [True, True],
    )
    self.assertEqual(
        (False | rasp.tokens)([True, False]),
        [True, False],
    )

    self.assertEqual(
        (~rasp.tokens)([True, False]),
        [False, True],
    )

    self.assertEqual(
        (rasp.ConstantSOp([True, True, False, False])
         & rasp.ConstantSOp([True, False, True, False]))([1, 1, 1, 1]),
        [True, False, False, False],
    )

    self.assertEqual(
        (rasp.ConstantSOp([True, True, False, False])
         | rasp.ConstantSOp([True, False, True, False]))([1, 1, 1, 1]),
        [True, True, True, False],
    )

class EncodingTest(parameterized.TestCase):
  """Tests for SOp encodings."""

  @parameterized.named_parameters(*_SOP_EXAMPLES())
  def test_all_sops_are_categorical_by_default(self, sop: rasp.SOp):
    self.assertTrue(rasp.is_categorical(sop))

  @parameterized.named_parameters(*_SOP_EXAMPLES())
  def test_is_numerical(self, sop: rasp.SOp):
    self.assertTrue(rasp.is_numerical(rasp.numerical(sop)))
    self.assertFalse(rasp.is_numerical(rasp.categorical(sop)))

  @parameterized.named_parameters(*_SOP_EXAMPLES())
  def test_is_categorical(self, sop: rasp.SOp):
    self.assertTrue(rasp.is_categorical(rasp.categorical(sop)))
    self.assertFalse(rasp.is_categorical(rasp.numerical(sop)))

  @parameterized.named_parameters(*_SOP_EXAMPLES())
  def test_double_encoding_annotations_overwrites_encoding(self, sop: rasp.SOp):
    num_sop = rasp.numerical(sop)
    cat_num_sop = rasp.categorical(num_sop)
    self.assertTrue(rasp.is_numerical(num_sop))
    self.assertTrue(rasp.is_categorical(cat_num_sop))

class SelectorTest(parameterized.TestCase):
  """Tests for Selectors."""

  def test_select_eq_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.EQ)
    self.assertEqual(
        selector("hey"), [
            [True, False, False],
            [False, True, False],
            [False, False, True],
        ])

  def test_select_lt_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.LT)
    self.assertEqual(selector([0, 1]), [
        [False, False],
        [True, False],
    ])

  def test_select_leq_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.LEQ)
    self.assertEqual(selector([0, 1]), [
        [True, False],
        [True, True],
    ])

  def test_select_gt_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.GT)
    self.assertEqual(selector([0, 1]), [
        [False, True],
        [False, False],
    ])

  def test_select_geq_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.GEQ)
    self.assertEqual(selector([0, 1]), [
        [True, True],
        [False, True],
    ])

  def test_select_neq_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.NEQ)
    self.assertEqual(selector([0, 1]), [
        [False, True],
        [True, False],
    ])

  def test_select_true_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.TRUE)
    self.assertEqual(selector([0, 1]), [
        [True, True],
        [True, True],
    ])

  def test_select_false_has_correct_value(self):
    selector = rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.FALSE)
    self.assertEqual(selector([0, 1]), [
        [False, False],
        [False, False],
    ])

  def test_selector_and_gets_simplified_when_keys_and_queries_match(self):
    selector = rasp.selector_and(
        rasp.Select(rasp.tokens, rasp.indices, rasp.Comparison.GEQ),
        rasp.Select(rasp.tokens, rasp.indices, rasp.Comparison.LEQ),
    )
    self.assertIsInstance(selector, rasp.Select)
    self.assertIs(selector.keys, rasp.tokens)
    self.assertIs(selector.queries, rasp.indices)

  def test_selector_and_doesnt_get_simplified_when_keys_queries_different(self):
    selector = rasp.selector_and(
        rasp.Select(rasp.tokens, rasp.indices, rasp.Comparison.GEQ),
        rasp.Select(rasp.indices, rasp.tokens, rasp.Comparison.LEQ),
    )
    self.assertIsInstance(selector, rasp.SelectorAnd)

  def test_selector_and_gets_simplified_when_keys_are_full(self):
    selector = rasp.selector_and(
        rasp.Select(rasp.Full(1), rasp.indices, rasp.Comparison.GEQ),
        rasp.Select(rasp.tokens, rasp.indices, rasp.Comparison.LEQ),
    )
    self.assertIsInstance(selector, rasp.Select)
    self.assertIs(selector.keys, rasp.tokens)
    self.assertIs(selector.queries, rasp.indices)

  def test_selector_and_gets_simplified_when_queries_are_full(self):
    selector = rasp.selector_and(
        rasp.Select(rasp.tokens, rasp.indices, rasp.Comparison.GEQ),
        rasp.Select(rasp.tokens, rasp.Full(1), rasp.Comparison.LEQ),
    )
    self.assertIsInstance(selector, rasp.Select)
    self.assertIs(selector.keys, rasp.tokens)
    self.assertIs(selector.queries, rasp.indices)

  @parameterized.parameters(
      itertools.product(
          (rasp.tokens, rasp.indices, rasp.Full(1)),
          (rasp.tokens, rasp.indices, rasp.Full(1)),
          list(rasp.Comparison),
          (rasp.tokens, rasp.indices, rasp.Full(1)),
          (rasp.tokens, rasp.indices, rasp.Full(1)),
          list(rasp.Comparison),
      ))
  def test_simplified_selector_and_works_the_same_way_as_not(
      self, fst_k, fst_q, fst_p, snd_k, snd_q, snd_p):
    fst = rasp.Select(fst_k, fst_q, fst_p)
    snd = rasp.Select(snd_k, snd_q, snd_p)

    simplified = rasp.selector_and(fst, snd)([0, 1, 2, 3])
    not_simplified = rasp.selector_and(fst, snd, simplify=False)([0, 1, 2, 3])

    np.testing.assert_array_equal(
        np.array(simplified),
        np.array(not_simplified),
    )

  def test_select_is_selector(self):
    self.assertIsInstance(
        rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.EQ),
        rasp.Selector,
    )

  def test_select_is_raspexpr(self):
    self.assertIsInstance(
        rasp.Select(rasp.tokens, rasp.tokens, rasp.Comparison.EQ),
        rasp.RASPExpr,
    )

  def test_constant_selector(self):
    self.assertEqual(
        rasp.ConstantSelector([[True, True], [False, False]])([1, 2]),
        [[True, True], [False, False]],
    )

class CopyTest(parameterized.TestCase):

  @parameterized.named_parameters(*_ALL_EXAMPLES())
  def test_copy_preserves_name(self, expr: rasp.RASPExpr):
    expr = expr.named("foo")
    self.assertEqual(expr.copy().name, expr.name)

  @parameterized.named_parameters(*_ALL_EXAMPLES())
  def test_renaming_copy_doesnt_rename_original(self, expr: rasp.RASPExpr):
    expr = expr.named("foo")
    expr.copy().named("bar")
    self.assertEqual(expr.name, "foo")

  @parameterized.named_parameters(*_ALL_EXAMPLES())
  def test_renaming_original_doesnt_rename_copy(self, expr: rasp.RASPExpr):
    expr = expr.named("foo")
    copy = expr.copy()
    expr.named("bar")
    self.assertEqual(copy.name, "foo")

  @parameterized.named_parameters(*_ALL_EXAMPLES())
  def test_copy_changes_id(self, expr: rasp.RASPExpr):
    self.assertNotEqual(expr.copy().unique_id, expr.unique_id)

  @parameterized.named_parameters(*_ALL_EXAMPLES())
  def test_copy_preserves_child_ids(self, expr: rasp.RASPExpr):
    copy_child_ids = [c.unique_id for c in expr.copy().children]
    child_ids = [c.unique_id for c in expr.children]
    for child_id, copy_child_id in zip(child_ids, copy_child_ids):
      self.assertEqual(child_id, copy_child_id)

class AggregateTest(parameterized.TestCase):
  """Tests for Aggregate."""

  @parameterized.parameters(
      dict(
          selector=rasp.ConstantSelector([
              [True, False],
              [False, True],
          ]),
          sop=rasp.ConstantSOp(["h", "e"]),
          default=None,
          expected_value=["h", "e"],
      ),
      dict(
          selector=rasp.ConstantSelector([
              [False, True],
              [False, False],
          ]),
          sop=rasp.ConstantSOp(["h", "e"]),
          default=None,
          expected_value=["e", None],
      ),
      dict(
          selector=rasp.ConstantSelector([
              [True, False],
              [False, False],
          ]),
          sop=rasp.ConstantSOp(["h", "e"]),
          default=None,
          expected_value=["h", None],
      ),
      dict(
          selector=rasp.ConstantSelector([
              [True, True],
              [False, True],
          ]),
          sop=rasp.ConstantSOp([0, 1]),
          default=0,
          expected_value=[0.5, 1],
      ),
      dict(
          selector=rasp.ConstantSelector([
              [False, False],
              [True, True],
          ]),
          sop=rasp.ConstantSOp([0, 1]),
          default=0,
          expected_value=[0, 0.5],
      ),
      dict(
          selector=rasp.ConstantSelector([
              [False, False],
              [True, True],
          ]),
          sop=rasp.ConstantSOp([0, 1]),
          default=None,
          expected_value=[None, 0.5],
      ),
  )
  def test_aggregate_on_size_2_inputs(self, selector, sop, default,
                                      expected_value):
    # The 0, 0 input is ignored as it's overridden by the constant SOps.
    self.assertEqual(
        rasp.Aggregate(selector, sop, default)([0, 0]),
        expected_value,
    )

class RaspProgramTest(parameterized.TestCase):
  """Each testcase implements and tests a RASP program."""

  def test_has_prev(self):

    def has_prev(seq: rasp.SOp) -> rasp.SOp:

      prev_copy = rasp.SelectorAnd(
          rasp.Select(seq, seq, rasp.Comparison.EQ),
          rasp.Select(rasp.indices, rasp.indices, rasp.Comparison.LT),
      )
      return rasp.Aggregate(prev_copy, rasp.Full(1), default=0) > 0


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Assemble weights of a transformer model from a craft residual stack."""

import dataclasses
from typing import Any, Callable, Optional, List, Tuple

import chex
import einops
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.craft import bases
from tracr.craft import transformers
from tracr.craft import vectorspace_fns
from tracr.transformer import encoder
from tracr.transformer import model
from typing_extensions import Protocol

@chex.dataclass
class AssembledTransformerModelOutput:
  decoded: List[Any]  # length T.
  unembedded: jax.Array  # [B, T]     B = 1 always.
  layer_outputs: List[jax.Array]  # [B, T, D]
  residuals: List[jax.Array]  # [B, T, D]
  attn_logits: List[jax.Array]  # [B, T, T, H]
  transformer_output: jax.Array  # [B, T, D]
  input_embeddings: jax.Array

class ModelForward(Protocol):

  def __call__(
      self,
      params: hk.Params,
      emb: jax.Array,
  ) -> model.CompiledTransformerModelOutput:
    """A hk-transformed forward pass through the compiled model."""

@dataclasses.dataclass
class AssembledTransformerModel:
  """Model architecture and parameters from assembling a model."""
  forward: ModelForward
  get_compiled_model: Callable[[], model.CompiledTransformerModel]
  params: hk.Params
  model_config: model.TransformerConfig
  residual_labels: List[str]
  input_encoder: Optional[encoder.Encoder] = None
  output_encoder: Optional[encoder.Encoder] = None

  def apply(self, tokens: List[bases.Value]) -> AssembledTransformerModelOutput:
    """Returns output from running the model on a set of input tokens."""

    if self.input_encoder:
      tokens = self.input_encoder.encode(tokens)
    tokens = jnp.array([tokens])
    output = self.forward(self.params, tokens)
    decoded = output.unembedded_output[0].tolist()
    if self.output_encoder:
      decoded = self.output_encoder.decode(decoded)

    if self.input_encoder.bos_token:
      # Special case for decoding the bos token position, for which the output
      # decoder might have unspecified behavior.
      decoded = [self.input_encoder.bos_token] + decoded[1:]

    return AssembledTransformerModelOutput(
        decoded=decoded,
        unembedded=output.unembedded_output,
        layer_outputs=output.transformer_output.layer_outputs,
        residuals=output.transformer_output.residuals,
        attn_logits=output.transformer_output.attn_logits,
        transformer_output=output.transformer_output.output,
        input_embeddings=output.transformer_output.input_embeddings)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Assemble weights of a transformer model from a craft residual stack."""

import dataclasses
from typing import Any, Callable, Optional, List, Tuple

import chex
import einops
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.craft import bases
from tracr.craft import transformers
from tracr.craft import vectorspace_fns
from tracr.transformer import encoder
from tracr.transformer import model
from typing_extensions import Protocol

@chex.dataclass
class AssembledTransformerModelOutput:
  decoded: List[Any]  # length T.
  unembedded: jax.Array  # [B, T]     B = 1 always.
  layer_outputs: List[jax.Array]  # [B, T, D]
  residuals: List[jax.Array]  # [B, T, D]
  attn_logits: List[jax.Array]  # [B, T, T, H]
  transformer_output: jax.Array  # [B, T, D]
  input_embeddings: jax.Array

class ModelForward(Protocol):

  def __call__(
      self,
      params: hk.Params,
      emb: jax.Array,
  ) -> model.CompiledTransformerModelOutput:
    """A hk-transformed forward pass through the compiled model."""

@dataclasses.dataclass
class AssembledTransformerModel:
  """Model architecture and parameters from assembling a model."""
  forward: ModelForward
  get_compiled_model: Callable[[], model.CompiledTransformerModel]
  params: hk.Params
  model_config: model.TransformerConfig
  residual_labels: List[str]
  input_encoder: Optional[encoder.Encoder] = None
  output_encoder: Optional[encoder.Encoder] = None

  def apply(self, tokens: List[bases.Value]) -> AssembledTransformerModelOutput:
    """Returns output from running the model on a set of input tokens."""
    if self.input_encoder:
      tokens = self.input_encoder.encode(tokens)
    tokens = jnp.array([tokens])
    output = self.forward(self.params, tokens)
    decoded = output.unembedded_output[0].tolist()
    if self.output_encoder:
      decoded = self.output_encoder.decode(decoded)

    if self.input_encoder.bos_token:
      # Special case for decoding the bos token position, for which the output
      # decoder might have unspecified behavior.
      decoded = [self.input_encoder.bos_token] + decoded[1:]

    return AssembledTransformerModelOutput(
        decoded=decoded,
        unembedded=output.unembedded_output,
        layer_outputs=output.transformer_output.layer_outputs,
        residuals=output.transformer_output.residuals,
        attn_logits=output.transformer_output.attn_logits,
        transformer_output=output.transformer_output.output,
        input_embeddings=output.transformer_output.input_embeddings)

@dataclasses.dataclass
class EmbeddingModules:
  """Modules for embedding and tokens and positions and unembedding results."""
  token_embed: model.CallableHaikuModule
  pos_embed: model.CallableHaikuModule
  unembed: model.CallableHaikuModule

def _get_model_config_and_module_names(
    craft_model: transformers.SeriesWithResiduals
) -> Tuple[model.TransformerConfig, List[str]]:
  """Returns model config and locations (in params) for halflayers."""

  multi_attn_heads: List[List[transformers.AttentionHead]] = []
  mlps: List[transformers.MLP] = []
  module_names: List[str] = []

  candidate_module_names = []
  for layer in range(len(craft_model.blocks)):
    candidate_module_names.append(f"transformer/layer_{layer}/attn")
    candidate_module_names.append(f"transformer/layer_{layer}/mlp")
  candidate_module_names = iter(candidate_module_names)

  for module in craft_model.blocks:
    if isinstance(module, transformers.MLP):
      mlps.append(module)
      layer_type = "mlp"
    else:
      multi_attn_heads.append(list(module.as_multi().heads()))
      layer_type = "attn"
    # Find next layer with the necessary type. Modules in-between, that are not
    # added to module_names will be disabled later by setting all weights to 0.
    module_name = next(candidate_module_names)
    while layer_type not in module_name:
      module_name = next(candidate_module_names)
    module_names.append(module_name)

  num_layers = int(module_names[-1].split("_")[1].split("/")[0]) + 1
  heads = sum(multi_attn_heads, [])

  if multi_attn_heads:
    num_heads = max(len(heads) for heads in multi_attn_heads)
    key_size = max(max(head.w_qk.matrix.shape) for head in heads)
  else:
    num_heads, key_size = 1, 1

  if mlps:
    mlp_hidden_size = max(mlp.fst.output_space.num_dims for mlp in mlps)
  else:
    mlp_hidden_size = 1

  model_config = model.TransformerConfig(
      num_heads=num_heads,
      num_layers=num_layers,
      key_size=key_size,
      mlp_hidden_size=mlp_hidden_size,
      dropout_rate=0.,
      activation_function=jax.nn.relu,
      layer_norm=False,
      causal=False,
  )

  return model_config, module_names

def _make_embedding_modules(
    residual_space: bases.VectorSpaceWithBasis,
    tokens_space: bases.VectorSpaceWithBasis,
    indices_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis) -> EmbeddingModules:
  """Creates embedding and unembedding modules from vector spaces.

  Args:
    residual_space: Full residual space of the model.
    tokens_space: Subspace to embed tokens to.
    indices_space: Subspace to embed indices/position embeddings to.
    output_space: Subspace to unembed outputs from.

  Returns:
    EmbeddingModules containing modules for token embeddings, position
      embeddings and unembeddings.
  """
  tokens_to_res = vectorspace_fns.project(tokens_space, residual_space)

  # If we use the 'one' direction, make sure all inputs have a 1 here
  one_dir = bases.BasisDirection("one")
  if one_dir in residual_space:
    one_to_res = vectorspace_fns.Linear.from_action(
        tokens_space, residual_space,
        lambda x: residual_space.vector_from_basis_direction(one_dir))
    tokens_to_res = vectorspace_fns.Linear.combine_in_parallel(
        [tokens_to_res, one_to_res])

  # Token embeddings.
  res_to_out = vectorspace_fns.project(residual_space, output_space)
  token_embed = hk.Embed(
      embedding_matrix=tokens_to_res.matrix, name="token_embed")

  # Positional embeddings.
  index_to_res = vectorspace_fns.project(indices_space, residual_space)
  # The zeroth position should not have any positional embeddings,
  # so we add one line of padding at the zeroth position.
  pos_matrix = np.concatenate(
      [np.zeros((1, residual_space.num_dims)), index_to_res.matrix], axis=0)
  pos_embed = hk.Embed(embedding_matrix=pos_matrix, name="pos_embed")

  def unembed(x, use_unembed_argmax):

    out = x @ res_to_out.matrix
    if use_unembed_argmax:
      return jnp.argmax(out, axis=-1)
    elif out.shape[-1] == 1:
      return out.squeeze(-1)
    return out


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Assemble weights of a transformer model from a craft residual stack."""

import dataclasses
from typing import Any, Callable, Optional, List, Tuple

import chex
import einops
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.craft import bases
from tracr.craft import transformers
from tracr.craft import vectorspace_fns
from tracr.transformer import encoder
from tracr.transformer import model
from typing_extensions import Protocol

@chex.dataclass
class AssembledTransformerModelOutput:
  decoded: List[Any]  # length T.
  unembedded: jax.Array  # [B, T]     B = 1 always.
  layer_outputs: List[jax.Array]  # [B, T, D]
  residuals: List[jax.Array]  # [B, T, D]
  attn_logits: List[jax.Array]  # [B, T, T, H]
  transformer_output: jax.Array  # [B, T, D]
  input_embeddings: jax.Array

class ModelForward(Protocol):

  def __call__(
      self,
      params: hk.Params,
      emb: jax.Array,
  ) -> model.CompiledTransformerModelOutput:
    """A hk-transformed forward pass through the compiled model."""

@dataclasses.dataclass
class AssembledTransformerModel:
  """Model architecture and parameters from assembling a model."""
  forward: ModelForward
  get_compiled_model: Callable[[], model.CompiledTransformerModel]
  params: hk.Params
  model_config: model.TransformerConfig
  residual_labels: List[str]
  input_encoder: Optional[encoder.Encoder] = None
  output_encoder: Optional[encoder.Encoder] = None

  def apply(self, tokens: List[bases.Value]) -> AssembledTransformerModelOutput:
    """Returns output from running the model on a set of input tokens."""
    if self.input_encoder:
      tokens = self.input_encoder.encode(tokens)
    tokens = jnp.array([tokens])
    output = self.forward(self.params, tokens)
    decoded = output.unembedded_output[0].tolist()
    if self.output_encoder:
      decoded = self.output_encoder.decode(decoded)

    if self.input_encoder.bos_token:
      # Special case for decoding the bos token position, for which the output
      # decoder might have unspecified behavior.
      decoded = [self.input_encoder.bos_token] + decoded[1:]

    return AssembledTransformerModelOutput(
        decoded=decoded,
        unembedded=output.unembedded_output,
        layer_outputs=output.transformer_output.layer_outputs,
        residuals=output.transformer_output.residuals,
        attn_logits=output.transformer_output.attn_logits,
        transformer_output=output.transformer_output.output,
        input_embeddings=output.transformer_output.input_embeddings)

@dataclasses.dataclass
class EmbeddingModules:
  """Modules for embedding and tokens and positions and unembedding results."""
  token_embed: model.CallableHaikuModule
  pos_embed: model.CallableHaikuModule
  unembed: model.CallableHaikuModule

def _get_model_config_and_module_names(
    craft_model: transformers.SeriesWithResiduals
) -> Tuple[model.TransformerConfig, List[str]]:
  """Returns model config and locations (in params) for halflayers."""

  multi_attn_heads: List[List[transformers.AttentionHead]] = []
  mlps: List[transformers.MLP] = []
  module_names: List[str] = []

  candidate_module_names = []
  for layer in range(len(craft_model.blocks)):
    candidate_module_names.append(f"transformer/layer_{layer}/attn")
    candidate_module_names.append(f"transformer/layer_{layer}/mlp")
  candidate_module_names = iter(candidate_module_names)

  for module in craft_model.blocks:
    if isinstance(module, transformers.MLP):
      mlps.append(module)
      layer_type = "mlp"
    else:
      multi_attn_heads.append(list(module.as_multi().heads()))
      layer_type = "attn"
    # Find next layer with the necessary type. Modules in-between, that are not
    # added to module_names will be disabled later by setting all weights to 0.
    module_name = next(candidate_module_names)
    while layer_type not in module_name:
      module_name = next(candidate_module_names)
    module_names.append(module_name)

  num_layers = int(module_names[-1].split("_")[1].split("/")[0]) + 1
  heads = sum(multi_attn_heads, [])

  if multi_attn_heads:
    num_heads = max(len(heads) for heads in multi_attn_heads)
    key_size = max(max(head.w_qk.matrix.shape) for head in heads)
  else:
    num_heads, key_size = 1, 1

  if mlps:
    mlp_hidden_size = max(mlp.fst.output_space.num_dims for mlp in mlps)
  else:
    mlp_hidden_size = 1

  model_config = model.TransformerConfig(
      num_heads=num_heads,
      num_layers=num_layers,
      key_size=key_size,
      mlp_hidden_size=mlp_hidden_size,
      dropout_rate=0.,
      activation_function=jax.nn.relu,
      layer_norm=False,
      causal=False,
  )

  return model_config, module_names

def _make_embedding_modules(
    residual_space: bases.VectorSpaceWithBasis,
    tokens_space: bases.VectorSpaceWithBasis,
    indices_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis) -> EmbeddingModules:
  """Creates embedding and unembedding modules from vector spaces.

  Args:
    residual_space: Full residual space of the model.
    tokens_space: Subspace to embed tokens to.
    indices_space: Subspace to embed indices/position embeddings to.
    output_space: Subspace to unembed outputs from.

  Returns:
    EmbeddingModules containing modules for token embeddings, position
      embeddings and unembeddings.
  """
  tokens_to_res = vectorspace_fns.project(tokens_space, residual_space)

  # If we use the 'one' direction, make sure all inputs have a 1 here
  one_dir = bases.BasisDirection("one")
  if one_dir in residual_space:
    one_to_res = vectorspace_fns.Linear.from_action(
        tokens_space, residual_space,
        lambda x: residual_space.vector_from_basis_direction(one_dir))
    tokens_to_res = vectorspace_fns.Linear.combine_in_parallel(
        [tokens_to_res, one_to_res])

  # Token embeddings.
  res_to_out = vectorspace_fns.project(residual_space, output_space)
  token_embed = hk.Embed(
      embedding_matrix=tokens_to_res.matrix, name="token_embed")

  # Positional embeddings.
  index_to_res = vectorspace_fns.project(indices_space, residual_space)
  # The zeroth position should not have any positional embeddings,
  # so we add one line of padding at the zeroth position.
  pos_matrix = np.concatenate(
      [np.zeros((1, residual_space.num_dims)), index_to_res.matrix], axis=0)
  pos_embed = hk.Embed(embedding_matrix=pos_matrix, name="pos_embed")

  def unembed(x, use_unembed_argmax):
    out = x @ res_to_out.matrix
    if use_unembed_argmax:
      return jnp.argmax(out, axis=-1)
    elif out.shape[-1] == 1:
      return out.squeeze(-1)
    return out

  unembed_mod = hk.to_module(unembed)()
  return EmbeddingModules(
      token_embed=token_embed, pos_embed=pos_embed, unembed=unembed_mod)

def assemble_craft_model(
    craft_model: transformers.SeriesWithResiduals,
    tokens_space: bases.VectorSpaceWithBasis,
    indices_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    categorical_output: bool,
    causal: bool = False,
) -> AssembledTransformerModel:
  """Assembles the given components into a Haiku model with parameters.

  Args:
    craft_model: Model to assemble weights for.
    tokens_space: Vectorspace to embed the input tokens to.
    indices_space: Vectorspace to embed the indices to (position encodings).
    output_space: Vectorspace that the model will write outputs to that should
      be unembedded.
    categorical_output: Whether the output is categorical. If True, we take an
      argmax when unembedding.
    causal: Whether to output a causally-masked model.

  Returns:
    An AssembledTransformerModel that contains the model and parameters of the
      assembled transformer.
  """
  # TODO(b/255936413): Make embeddings only retain the tokens and indices that
  #   are actually used.
  # TODO(b/255936496): Think about enabling layer norm and reversing it somehow

  model_config, module_names = _get_model_config_and_module_names(craft_model)
  model_config.causal = causal

  residual_space = bases.join_vector_spaces(craft_model.residual_space,
                                            tokens_space, indices_space,
                                            output_space)
  residual_labels = [str(basis_dir) for basis_dir in residual_space.basis]

  # Build model with embedding and unembedding layers
  def get_compiled_model():

    transformer = model.Transformer(model_config)
    embed_modules = _make_embedding_modules(
        residual_space=residual_space,
        tokens_space=tokens_space,
        indices_space=indices_space,
        output_space=output_space)
    return model.CompiledTransformerModel(
        transformer=transformer,
        token_embed=embed_modules.token_embed,
        position_embed=embed_modules.pos_embed,
        unembed=embed_modules.unembed,
        use_unembed_argmax=categorical_output)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Integration tests for the full RASP -> transformer compilation."""

from absl.testing import absltest
from absl.testing import parameterized
import jax
import numpy as np

from tracr.compiler import compiling
from tracr.compiler import lib
from tracr.compiler import test_cases
from tracr.craft import tests_common
from tracr.rasp import rasp

_COMPILER_BOS = "rasp_to_transformer_integration_test_BOS"
_COMPILER_PAD = "rasp_to_transformer_integration_test_PAD"

# Force float32 precision on TPU, which otherwise defaults to float16.
jax.config.update("jax_default_matmul_precision", "float32")

class CompilerIntegrationTest(tests_common.VectorFnTestCase):

  def assertSequenceEqualWhenExpectedIsNotNone(self, actual_seq, expected_seq):

    for actual, expected in zip(actual_seq, expected_seq):
      if expected is not None and actual != expected:
        self.fail(f"{actual_seq} does not match (ignoring Nones) "
                  f"expected_seq={expected_seq}")


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for compiler.expr_to_craft_graph."""

from absl.testing import absltest
from absl.testing import parameterized
from tracr.compiler import basis_inference
from tracr.compiler import expr_to_craft_graph
from tracr.compiler import lib
from tracr.compiler import nodes
from tracr.compiler import rasp_to_graph
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

class ExprToCraftGraphTest(parameterized.TestCase):

  def _check_block_types_are_correct(self, graph):

    for _, node in graph.nodes.items():
      expr = node[nodes.EXPR]
      if isinstance(expr, rasp.SOp):
        block = node[nodes.MODEL_BLOCK]
        if isinstance(expr, (rasp.Map, rasp.SequenceMap)):
          self.assertIsInstance(block, transformers.MLP)
        elif isinstance(expr, rasp.Aggregate):
          self.assertIsInstance(block, transformers.AttentionHead)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for compiler.expr_to_craft_graph."""

from absl.testing import absltest
from absl.testing import parameterized
from tracr.compiler import basis_inference
from tracr.compiler import expr_to_craft_graph
from tracr.compiler import lib
from tracr.compiler import nodes
from tracr.compiler import rasp_to_graph
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

class ExprToCraftGraphTest(parameterized.TestCase):

  def _check_block_types_are_correct(self, graph):
    for _, node in graph.nodes.items():
      expr = node[nodes.EXPR]
      if isinstance(expr, rasp.SOp):
        block = node[nodes.MODEL_BLOCK]
        if isinstance(expr, (rasp.Map, rasp.SequenceMap)):
          self.assertIsInstance(block, transformers.MLP)
        elif isinstance(expr, rasp.Aggregate):
          self.assertIsInstance(block, transformers.AttentionHead)

  def _get_input_space_from_node(self, node):

    block = node[nodes.MODEL_BLOCK]
    if isinstance(block, transformers.MLP):
      return block.fst.input_space
    elif isinstance(block, transformers.AttentionHead):
      return bases.join_vector_spaces(block.w_qk.left_space,
                                      block.w_qk.right_space,
                                      block.w_ov.input_space)
    else:
      return None


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for compiler.expr_to_craft_graph."""

from absl.testing import absltest
from absl.testing import parameterized
from tracr.compiler import basis_inference
from tracr.compiler import expr_to_craft_graph
from tracr.compiler import lib
from tracr.compiler import nodes
from tracr.compiler import rasp_to_graph
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

class ExprToCraftGraphTest(parameterized.TestCase):

  def _check_block_types_are_correct(self, graph):
    for _, node in graph.nodes.items():
      expr = node[nodes.EXPR]
      if isinstance(expr, rasp.SOp):
        block = node[nodes.MODEL_BLOCK]
        if isinstance(expr, (rasp.Map, rasp.SequenceMap)):
          self.assertIsInstance(block, transformers.MLP)
        elif isinstance(expr, rasp.Aggregate):
          self.assertIsInstance(block, transformers.AttentionHead)

  def _get_input_space_from_node(self, node):
    block = node[nodes.MODEL_BLOCK]
    if isinstance(block, transformers.MLP):
      return block.fst.input_space
    elif isinstance(block, transformers.AttentionHead):
      return bases.join_vector_spaces(block.w_qk.left_space,
                                      block.w_qk.right_space,
                                      block.w_ov.input_space)
    else:
      return None

  def _check_spaces_are_consistent(self, graph):
    """Check that for each edge the output is a subspace of the input."""

    for u, v in graph.edges:
      u_node, v_node = graph.nodes[u], graph.nodes[v]
      if isinstance(u_node[nodes.EXPR], rasp.SOp) and isinstance(
          v_node[nodes.EXPR], rasp.SOp):
        u_out_basis = u_node[nodes.OUTPUT_BASIS]
        u_out_space = bases.VectorSpaceWithBasis(u_out_basis)
        v_in_space = self._get_input_space_from_node(v_node)
        self.assertTrue(u_out_space.issubspace(v_in_space))


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a craft model from a computational graph."""

import collections
from typing import Dict, List, Sequence

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

def _get_longest_path_length_to_node(graph: nx.DiGraph, sources: Sequence[Node],
                                     node: Node) -> int:
  """Returns the lengths of the longest path from sources to node.

  Only SOps count towards the length of a path.

  Args:
    graph: DAG to compute longest path in.
    sources: List of starting nodes, longest path will be a maximum over all.
    node: Target node.

  Returns:
    Number of steps needed for the longest path from the source to the node, or
    -1 if there is no path from any of the sources to the target node.
  """

  if node in sources:
    return 0

  def num_sops(path: Sequence[NodeID]) -> int:
    num = 0
    for node_id in path:
      if isinstance(graph.nodes[node_id][nodes.EXPR], rasp.SOp):
        num += 1
    return num

  result = -1
  for source in sources:
    all_paths = nx.all_simple_paths(graph, source[nodes.ID], node[nodes.ID])
    longest_path_len = max(map(num_sops, all_paths), default=-1) - 1
    if longest_path_len > result:
      result = longest_path_len
  return result


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a craft model from a computational graph."""

import collections
from typing import Dict, List, Sequence

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

def _get_longest_path_length_to_node(graph: nx.DiGraph, sources: Sequence[Node],
                                     node: Node) -> int:
  """Returns the lengths of the longest path from sources to node.

  Only SOps count towards the length of a path.

  Args:
    graph: DAG to compute longest path in.
    sources: List of starting nodes, longest path will be a maximum over all.
    node: Target node.

  Returns:
    Number of steps needed for the longest path from the source to the node, or
    -1 if there is no path from any of the sources to the target node.
  """
  if node in sources:
    return 0

  def num_sops(path: Sequence[NodeID]) -> int:

    num = 0
    for node_id in path:
      if isinstance(graph.nodes[node_id][nodes.EXPR], rasp.SOp):
        num += 1
    return num


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a craft model from a computational graph."""

import collections
from typing import Dict, List, Sequence

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

def _get_longest_path_length_to_node(graph: nx.DiGraph, sources: Sequence[Node],
                                     node: Node) -> int:
  """Returns the lengths of the longest path from sources to node.

  Only SOps count towards the length of a path.

  Args:
    graph: DAG to compute longest path in.
    sources: List of starting nodes, longest path will be a maximum over all.
    node: Target node.

  Returns:
    Number of steps needed for the longest path from the source to the node, or
    -1 if there is no path from any of the sources to the target node.
  """
  if node in sources:
    return 0

  def num_sops(path: Sequence[NodeID]) -> int:
    num = 0
    for node_id in path:
      if isinstance(graph.nodes[node_id][nodes.EXPR], rasp.SOp):
        num += 1
    return num

  result = -1
  for source in sources:
    all_paths = nx.all_simple_paths(graph, source[nodes.ID], node[nodes.ID])
    longest_path_len = max(map(num_sops, all_paths), default=-1) - 1
    if longest_path_len > result:
      result = longest_path_len
  return result

def _node_is_attn(node: Node) -> bool:
  """Returns True if node is an attention layer."""

  return nodes.MODEL_BLOCK in node and isinstance(
      node[nodes.MODEL_BLOCK],
      (transformers.AttentionHead, transformers.MultiAttentionHead))


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a craft model from a computational graph."""

import collections
from typing import Dict, List, Sequence

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

def _get_longest_path_length_to_node(graph: nx.DiGraph, sources: Sequence[Node],
                                     node: Node) -> int:
  """Returns the lengths of the longest path from sources to node.

  Only SOps count towards the length of a path.

  Args:
    graph: DAG to compute longest path in.
    sources: List of starting nodes, longest path will be a maximum over all.
    node: Target node.

  Returns:
    Number of steps needed for the longest path from the source to the node, or
    -1 if there is no path from any of the sources to the target node.
  """
  if node in sources:
    return 0

  def num_sops(path: Sequence[NodeID]) -> int:
    num = 0
    for node_id in path:
      if isinstance(graph.nodes[node_id][nodes.EXPR], rasp.SOp):
        num += 1
    return num

  result = -1
  for source in sources:
    all_paths = nx.all_simple_paths(graph, source[nodes.ID], node[nodes.ID])
    longest_path_len = max(map(num_sops, all_paths), default=-1) - 1
    if longest_path_len > result:
      result = longest_path_len
  return result

def _node_is_attn(node: Node) -> bool:
  """Returns True if node is an attention layer."""
  return nodes.MODEL_BLOCK in node and isinstance(
      node[nodes.MODEL_BLOCK],
      (transformers.AttentionHead, transformers.MultiAttentionHead))

def _node_is_mlp(node: Node) -> bool:
  """Returns True if node is an MLP layer."""
  return nodes.MODEL_BLOCK in node and isinstance(node[nodes.MODEL_BLOCK],
                                                  transformers.MLP)

def _node_is_residual_block(node: Node) -> bool:
  """Returns True if node is a valid residual block (Attn followed by MLP)."""

  block = node[nodes.MODEL_BLOCK] if nodes.MODEL_BLOCK in node else None
  if block and isinstance(block, transformers.SeriesWithResiduals):
    if len(block.blocks) == 2:
      attn, mlp = block.blocks
      if (isinstance(
          attn,
          (transformers.AttentionHead, transformers.MultiAttentionHead)) and
          isinstance(mlp, transformers.MLP)):
        return True
  return False


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a craft model from a computational graph."""

import collections
from typing import Dict, List, Sequence

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

def _get_longest_path_length_to_node(graph: nx.DiGraph, sources: Sequence[Node],
                                     node: Node) -> int:
  """Returns the lengths of the longest path from sources to node.

  Only SOps count towards the length of a path.

  Args:
    graph: DAG to compute longest path in.
    sources: List of starting nodes, longest path will be a maximum over all.
    node: Target node.

  Returns:
    Number of steps needed for the longest path from the source to the node, or
    -1 if there is no path from any of the sources to the target node.
  """
  if node in sources:
    return 0

  def num_sops(path: Sequence[NodeID]) -> int:
    num = 0
    for node_id in path:
      if isinstance(graph.nodes[node_id][nodes.EXPR], rasp.SOp):
        num += 1
    return num

  result = -1
  for source in sources:
    all_paths = nx.all_simple_paths(graph, source[nodes.ID], node[nodes.ID])
    longest_path_len = max(map(num_sops, all_paths), default=-1) - 1
    if longest_path_len > result:
      result = longest_path_len
  return result

def _node_is_attn(node: Node) -> bool:
  """Returns True if node is an attention layer."""
  return nodes.MODEL_BLOCK in node and isinstance(
      node[nodes.MODEL_BLOCK],
      (transformers.AttentionHead, transformers.MultiAttentionHead))

def _node_is_mlp(node: Node) -> bool:
  """Returns True if node is an MLP layer."""
  return nodes.MODEL_BLOCK in node and isinstance(node[nodes.MODEL_BLOCK],
                                                  transformers.MLP)

def _node_is_residual_block(node: Node) -> bool:
  """Returns True if node is a valid residual block (Attn followed by MLP)."""
  block = node[nodes.MODEL_BLOCK] if nodes.MODEL_BLOCK in node else None
  if block and isinstance(block, transformers.SeriesWithResiduals):
    if len(block.blocks) == 2:
      attn, mlp = block.blocks
      if (isinstance(
          attn,
          (transformers.AttentionHead, transformers.MultiAttentionHead)) and
          isinstance(mlp, transformers.MLP)):
        return True
  return False

def _all_attn_nodes(node_list: Sequence[Node]) -> bool:
  """Returns True iff all nodes are attention layers (or nodes is empty)."""

  for node in node_list:
    if not _node_is_attn(node):
      return False
  return True


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Create a craft model from a computational graph."""

import collections
from typing import Dict, List, Sequence

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft import transformers
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

def _get_longest_path_length_to_node(graph: nx.DiGraph, sources: Sequence[Node],
                                     node: Node) -> int:
  """Returns the lengths of the longest path from sources to node.

  Only SOps count towards the length of a path.

  Args:
    graph: DAG to compute longest path in.
    sources: List of starting nodes, longest path will be a maximum over all.
    node: Target node.

  Returns:
    Number of steps needed for the longest path from the source to the node, or
    -1 if there is no path from any of the sources to the target node.
  """
  if node in sources:
    return 0

  def num_sops(path: Sequence[NodeID]) -> int:
    num = 0
    for node_id in path:
      if isinstance(graph.nodes[node_id][nodes.EXPR], rasp.SOp):
        num += 1
    return num

  result = -1
  for source in sources:
    all_paths = nx.all_simple_paths(graph, source[nodes.ID], node[nodes.ID])
    longest_path_len = max(map(num_sops, all_paths), default=-1) - 1
    if longest_path_len > result:
      result = longest_path_len
  return result

def _node_is_attn(node: Node) -> bool:
  """Returns True if node is an attention layer."""
  return nodes.MODEL_BLOCK in node and isinstance(
      node[nodes.MODEL_BLOCK],
      (transformers.AttentionHead, transformers.MultiAttentionHead))

def _node_is_mlp(node: Node) -> bool:
  """Returns True if node is an MLP layer."""
  return nodes.MODEL_BLOCK in node and isinstance(node[nodes.MODEL_BLOCK],
                                                  transformers.MLP)

def _node_is_residual_block(node: Node) -> bool:
  """Returns True if node is a valid residual block (Attn followed by MLP)."""
  block = node[nodes.MODEL_BLOCK] if nodes.MODEL_BLOCK in node else None
  if block and isinstance(block, transformers.SeriesWithResiduals):
    if len(block.blocks) == 2:
      attn, mlp = block.blocks
      if (isinstance(
          attn,
          (transformers.AttentionHead, transformers.MultiAttentionHead)) and
          isinstance(mlp, transformers.MLP)):
        return True
  return False

def _all_attn_nodes(node_list: Sequence[Node]) -> bool:
  """Returns True iff all nodes are attention layers (or nodes is empty)."""
  for node in node_list:
    if not _node_is_attn(node):
      return False
  return True

def _all_mlp_nodes(node_list: Sequence[Node]) -> bool:
  """Returns True iff all nodes are MLP layers (or nodes is empty)."""

  for node in node_list:
    if not _node_is_mlp(node):
      return False
  return True


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Add craft model blocks to graph of RASPExpr."""

from typing import Any, Callable, Optional

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft.chamber import categorical_attn
from tracr.craft.chamber import categorical_mlp
from tracr.craft.chamber import numerical_mlp
from tracr.craft.chamber import selector_width
from tracr.rasp import rasp

def _transform_fun_to_basis_fun(
    fun: Callable[..., Any],
    output_direction_name: Optional[str] = None) -> Callable[..., Any]:
  """Transforms a function acting on values into one acting on directions."""

  def bases_fun(*args):
    values = [d.value for d in args]
    result = fun(*values)
    if output_direction_name:
      return bases.BasisDirection(output_direction_name, result)
    return result

  return bases_fun


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Add craft model blocks to graph of RASPExpr."""

from typing import Any, Callable, Optional

import networkx as nx
from tracr.compiler import nodes
from tracr.craft import bases
from tracr.craft.chamber import categorical_attn
from tracr.craft.chamber import categorical_mlp
from tracr.craft.chamber import numerical_mlp
from tracr.craft.chamber import selector_width
from tracr.rasp import rasp

def _transform_fun_to_basis_fun(
    fun: Callable[..., Any],
    output_direction_name: Optional[str] = None) -> Callable[..., Any]:
  """Transforms a function acting on values into one acting on directions."""

  def bases_fun(*args):

    values = [d.value for d in args]
    result = fun(*values)
    if output_direction_name:
      return bases.BasisDirection(output_direction_name, result)
    return result


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for compiler.craft_graph_to_model."""

from absl.testing import absltest
from absl.testing import parameterized
import networkx as nx
from tracr.compiler import craft_graph_to_model
from tracr.compiler import nodes
from tracr.compiler import rasp_to_graph
from tracr.craft import bases
from tracr.craft.chamber import categorical_attn
from tracr.craft.chamber import categorical_mlp
from tracr.rasp import rasp

class CraftAllocateModulesToLayersTest(parameterized.TestCase):

  def _get_dummy_block(self, block_type):

    if block_type == "ATTN":
      return categorical_attn.categorical_attn(
          query_space=bases.VectorSpaceWithBasis.from_names(["query"]),
          key_space=bases.VectorSpaceWithBasis.from_names(["bos", "key"]),
          value_space=bases.VectorSpaceWithBasis.from_names(["bos", "value"]),
          output_space=bases.VectorSpaceWithBasis.from_names(["output"]),
          bos_space=bases.VectorSpaceWithBasis.from_names(["bos"]),
          one_space=bases.VectorSpaceWithBasis.from_names(["one"]),
          attn_fn=lambda x, y: True,
      )
    elif block_type == "MLP":
      return categorical_mlp.map_categorical_mlp(
          input_space=bases.VectorSpaceWithBasis.from_names(["input"]),
          output_space=bases.VectorSpaceWithBasis.from_names(["output"]),
          operation=lambda x: x,
      )
    else:
      return None


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converting a RaspExpr to a graph."""

import dataclasses
import queue
from typing import List

import networkx as nx
from tracr.compiler import nodes
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

@dataclasses.dataclass
class ExtractRaspGraphOutput:
  graph: nx.DiGraph
  sink: Node  # the program's output.
  sources: List[Node]  # the primitive S-Ops.

def extract_rasp_graph(tip: rasp.SOp) -> ExtractRaspGraphOutput:
  """Converts a RASP program into a graph representation."""
  expr_queue = queue.Queue()
  graph = nx.DiGraph()
  sources: List[NodeID] = []

  def ensure_node(expr: rasp.RASPExpr) -> NodeID:
    """Finds or creates a graph node corresponding to expr; returns its ID."""

    node_id = expr.label
    if node_id not in graph:
      graph.add_node(node_id, **{nodes.ID: node_id, nodes.EXPR: expr})

    return node_id


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Converting a RaspExpr to a graph."""

import dataclasses
import queue
from typing import List

import networkx as nx
from tracr.compiler import nodes
from tracr.rasp import rasp

Node = nodes.Node
NodeID = nodes.NodeID

@dataclasses.dataclass
class ExtractRaspGraphOutput:
  graph: nx.DiGraph
  sink: Node  # the program's output.
  sources: List[Node]  # the primitive S-Ops.

def extract_rasp_graph(tip: rasp.SOp) -> ExtractRaspGraphOutput:
  """Converts a RASP program into a graph representation."""
  expr_queue = queue.Queue()
  graph = nx.DiGraph()
  sources: List[NodeID] = []

  def ensure_node(expr: rasp.RASPExpr) -> NodeID:
    """Finds or creates a graph node corresponding to expr; returns its ID."""
    node_id = expr.label
    if node_id not in graph:
      graph.add_node(node_id, **{nodes.ID: node_id, nodes.EXPR: expr})

    return node_id

  # Breadth-first search over the RASP expression graph.

  def visit_raspexpr(expr: rasp.RASPExpr):

    parent_id = ensure_node(expr)

    for child_expr in expr.children:
      expr_queue.put(child_expr)
      child_id = ensure_node(child_expr)
      graph.add_edge(child_id, parent_id)

    if not expr.children:
      sources.append(graph.nodes[parent_id])


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Integration tests for the RASP -> craft stages of the compiler."""

import unittest

from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from tracr.compiler import basis_inference
from tracr.compiler import craft_graph_to_model
from tracr.compiler import expr_to_craft_graph
from tracr.compiler import nodes
from tracr.compiler import rasp_to_graph
from tracr.compiler import test_cases
from tracr.craft import bases
from tracr.craft import tests_common
from tracr.rasp import rasp

_BOS_DIRECTION = "rasp_to_transformer_integration_test_BOS"
_ONE_DIRECTION = "rasp_to_craft_integration_test_ONE"

def _make_input_space(vocab, max_seq_len):

  tokens_space = bases.VectorSpaceWithBasis.from_values("tokens", vocab)
  indices_space = bases.VectorSpaceWithBasis.from_values(
      "indices", range(max_seq_len))
  one_space = bases.VectorSpaceWithBasis.from_names([_ONE_DIRECTION])
  bos_space = bases.VectorSpaceWithBasis.from_names([_BOS_DIRECTION])
  input_space = bases.join_vector_spaces(tokens_space, indices_space, one_space,
                                         bos_space)

  return input_space


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Integration tests for the RASP -> craft stages of the compiler."""

import unittest

from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from tracr.compiler import basis_inference
from tracr.compiler import craft_graph_to_model
from tracr.compiler import expr_to_craft_graph
from tracr.compiler import nodes
from tracr.compiler import rasp_to_graph
from tracr.compiler import test_cases
from tracr.craft import bases
from tracr.craft import tests_common
from tracr.rasp import rasp

_BOS_DIRECTION = "rasp_to_transformer_integration_test_BOS"
_ONE_DIRECTION = "rasp_to_craft_integration_test_ONE"

def _make_input_space(vocab, max_seq_len):
  tokens_space = bases.VectorSpaceWithBasis.from_values("tokens", vocab)
  indices_space = bases.VectorSpaceWithBasis.from_values(
      "indices", range(max_seq_len))
  one_space = bases.VectorSpaceWithBasis.from_names([_ONE_DIRECTION])
  bos_space = bases.VectorSpaceWithBasis.from_names([_BOS_DIRECTION])
  input_space = bases.join_vector_spaces(tokens_space, indices_space, one_space,
                                         bos_space)

  return input_space

def _embed_input(input_seq, input_space):

  bos_vec = input_space.vector_from_basis_direction(
      bases.BasisDirection(_BOS_DIRECTION))
  one_vec = input_space.vector_from_basis_direction(
      bases.BasisDirection(_ONE_DIRECTION))
  embedded_input = [bos_vec + one_vec]
  for i, val in enumerate(input_seq):
    i_vec = input_space.vector_from_basis_direction(
        bases.BasisDirection("indices", i))
    val_vec = input_space.vector_from_basis_direction(
        bases.BasisDirection("tokens", val))
    embedded_input.append(i_vec + val_vec + one_vec)
  return bases.VectorInBasis.stack(embedded_input)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Integration tests for the RASP -> craft stages of the compiler."""

import unittest

from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from tracr.compiler import basis_inference
from tracr.compiler import craft_graph_to_model
from tracr.compiler import expr_to_craft_graph
from tracr.compiler import nodes
from tracr.compiler import rasp_to_graph
from tracr.compiler import test_cases
from tracr.craft import bases
from tracr.craft import tests_common
from tracr.rasp import rasp

_BOS_DIRECTION = "rasp_to_transformer_integration_test_BOS"
_ONE_DIRECTION = "rasp_to_craft_integration_test_ONE"

def _make_input_space(vocab, max_seq_len):
  tokens_space = bases.VectorSpaceWithBasis.from_values("tokens", vocab)
  indices_space = bases.VectorSpaceWithBasis.from_values(
      "indices", range(max_seq_len))
  one_space = bases.VectorSpaceWithBasis.from_names([_ONE_DIRECTION])
  bos_space = bases.VectorSpaceWithBasis.from_names([_BOS_DIRECTION])
  input_space = bases.join_vector_spaces(tokens_space, indices_space, one_space,
                                         bos_space)

  return input_space

def _embed_input(input_seq, input_space):
  bos_vec = input_space.vector_from_basis_direction(
      bases.BasisDirection(_BOS_DIRECTION))
  one_vec = input_space.vector_from_basis_direction(
      bases.BasisDirection(_ONE_DIRECTION))
  embedded_input = [bos_vec + one_vec]
  for i, val in enumerate(input_seq):
    i_vec = input_space.vector_from_basis_direction(
        bases.BasisDirection("indices", i))
    val_vec = input_space.vector_from_basis_direction(
        bases.BasisDirection("tokens", val))
    embedded_input.append(i_vec + val_vec + one_vec)
  return bases.VectorInBasis.stack(embedded_input)

def _embed_output(output_seq, output_space, categorical_output):

  embedded_output = []
  output_label = output_space.basis[0].name
  for x in output_seq:
    if x is None:
      out_vec = output_space.null_vector()
    elif categorical_output:
      out_vec = output_space.vector_from_basis_direction(
          bases.BasisDirection(output_label, x))
    else:
      out_vec = x * output_space.vector_from_basis_direction(
          output_space.basis[0])
    embedded_output.append(out_vec)
  return bases.VectorInBasis.stack(embedded_output)


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP programs only using the subset of RASP supported by the compiler."""

from typing import List, Sequence

from tracr.rasp import rasp

### Programs that work only under non-causal evaluation.

def make_length() -> rasp.SOp:
  """Creates the `length` SOp using selector width primitive.

  Example usage:
    length = make_length()
    length("abcdefg")
    >> [7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0]

  Returns:
    length: SOp mapping an input to a sequence, where every element
      is the length of that sequence.
  """

  all_true_selector = rasp.Select(
      rasp.tokens, rasp.tokens, rasp.Comparison.TRUE).named("all_true_selector")
  return rasp.SelectorWidth(all_true_selector).named("length")


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP programs only using the subset of RASP supported by the compiler."""

from typing import List, Sequence

from tracr.rasp import rasp

### Programs that work only under non-causal evaluation.

def make_length() -> rasp.SOp:
  """Creates the `length` SOp using selector width primitive.

  Example usage:
    length = make_length()
    length("abcdefg")
    >> [7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0]

  Returns:
    length: SOp mapping an input to a sequence, where every element
      is the length of that sequence.
  """
  all_true_selector = rasp.Select(
      rasp.tokens, rasp.tokens, rasp.Comparison.TRUE).named("all_true_selector")
  return rasp.SelectorWidth(all_true_selector).named("length")

length = make_length()

def make_reverse(sop: rasp.SOp) -> rasp.SOp:
  """Create an SOp that reverses a sequence, using length primitive.

  Example usage:
    reverse = make_reverse(rasp.tokens)
    reverse("Hello")
    >> ['o', 'l', 'l', 'e', 'H']

  Args:
    sop: an SOp

  Returns:
    reverse : SOp that reverses the input sequence.
  """

  opp_idx = (length - rasp.indices).named("opp_idx")
  opp_idx = (opp_idx - 1).named("opp_idx-1")
  reverse_selector = rasp.Select(rasp.indices, opp_idx,
                                 rasp.Comparison.EQ).named("reverse_selector")
  return rasp.Aggregate(reverse_selector, sop).named("reverse")


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP programs only using the subset of RASP supported by the compiler."""

from typing import List, Sequence

from tracr.rasp import rasp

### Programs that work only under non-causal evaluation.

def make_length() -> rasp.SOp:
  """Creates the `length` SOp using selector width primitive.

  Example usage:
    length = make_length()
    length("abcdefg")
    >> [7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0]

  Returns:
    length: SOp mapping an input to a sequence, where every element
      is the length of that sequence.
  """
  all_true_selector = rasp.Select(
      rasp.tokens, rasp.tokens, rasp.Comparison.TRUE).named("all_true_selector")
  return rasp.SelectorWidth(all_true_selector).named("length")

length = make_length()

def make_reverse(sop: rasp.SOp) -> rasp.SOp:
  """Create an SOp that reverses a sequence, using length primitive.

  Example usage:
    reverse = make_reverse(rasp.tokens)
    reverse("Hello")
    >> ['o', 'l', 'l', 'e', 'H']

  Args:
    sop: an SOp

  Returns:
    reverse : SOp that reverses the input sequence.
  """
  opp_idx = (length - rasp.indices).named("opp_idx")
  opp_idx = (opp_idx - 1).named("opp_idx-1")
  reverse_selector = rasp.Select(rasp.indices, opp_idx,
                                 rasp.Comparison.EQ).named("reverse_selector")
  return rasp.Aggregate(reverse_selector, sop).named("reverse")

def make_pair_balance(sop: rasp.SOp, open_token: str,
                      close_token: str) -> rasp.SOp:
  """Return fraction of previous open tokens minus the fraction of close tokens.

   (As implemented in the RASP paper.)

  If the outputs are always non-negative and end in 0, that implies the input
  has balanced parentheses.

  Example usage:
    num_l = make_pair_balance(rasp.tokens, "(", ")")
    num_l("a()b(c))")
    >> [0, 1/2, 0, 0, 1/5, 1/6, 0, -1/8]

  Args:
    sop: Input SOp.
    open_token: Token that counts positive.
    close_token: Token that counts negative.

  Returns:
    pair_balance: SOp mapping an input to a sequence, where every element
      is the fraction of previous open tokens minus previous close tokens.
  """

  bools_open = rasp.numerical(sop == open_token).named("bools_open")
  opens = rasp.numerical(make_frac_prevs(bools_open)).named("opens")

  bools_close = rasp.numerical(sop == close_token).named("bools_close")
  closes = rasp.numerical(make_frac_prevs(bools_close)).named("closes")

  pair_balance = rasp.numerical(rasp.LinearSequenceMap(opens, closes, 1, -1))
  return pair_balance.named("pair_balance")


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP programs only using the subset of RASP supported by the compiler."""

from typing import List, Sequence

from tracr.rasp import rasp

### Programs that work only under non-causal evaluation.

def make_length() -> rasp.SOp:
  """Creates the `length` SOp using selector width primitive.

  Example usage:
    length = make_length()
    length("abcdefg")
    >> [7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0]

  Returns:
    length: SOp mapping an input to a sequence, where every element
      is the length of that sequence.
  """
  all_true_selector = rasp.Select(
      rasp.tokens, rasp.tokens, rasp.Comparison.TRUE).named("all_true_selector")
  return rasp.SelectorWidth(all_true_selector).named("length")

length = make_length()

def make_reverse(sop: rasp.SOp) -> rasp.SOp:
  """Create an SOp that reverses a sequence, using length primitive.

  Example usage:
    reverse = make_reverse(rasp.tokens)
    reverse("Hello")
    >> ['o', 'l', 'l', 'e', 'H']

  Args:
    sop: an SOp

  Returns:
    reverse : SOp that reverses the input sequence.
  """
  opp_idx = (length - rasp.indices).named("opp_idx")
  opp_idx = (opp_idx - 1).named("opp_idx-1")
  reverse_selector = rasp.Select(rasp.indices, opp_idx,
                                 rasp.Comparison.EQ).named("reverse_selector")
  return rasp.Aggregate(reverse_selector, sop).named("reverse")

def make_pair_balance(sop: rasp.SOp, open_token: str,
                      close_token: str) -> rasp.SOp:
  """Return fraction of previous open tokens minus the fraction of close tokens.

   (As implemented in the RASP paper.)

  If the outputs are always non-negative and end in 0, that implies the input
  has balanced parentheses.

  Example usage:
    num_l = make_pair_balance(rasp.tokens, "(", ")")
    num_l("a()b(c))")
    >> [0, 1/2, 0, 0, 1/5, 1/6, 0, -1/8]

  Args:
    sop: Input SOp.
    open_token: Token that counts positive.
    close_token: Token that counts negative.

  Returns:
    pair_balance: SOp mapping an input to a sequence, where every element
      is the fraction of previous open tokens minus previous close tokens.
  """
  bools_open = rasp.numerical(sop == open_token).named("bools_open")
  opens = rasp.numerical(make_frac_prevs(bools_open)).named("opens")

  bools_close = rasp.numerical(sop == close_token).named("bools_close")
  closes = rasp.numerical(make_frac_prevs(bools_close)).named("closes")

  pair_balance = rasp.numerical(rasp.LinearSequenceMap(opens, closes, 1, -1))
  return pair_balance.named("pair_balance")

def make_shuffle_dyck(pairs: List[str]) -> rasp.SOp:
  """Returns 1 if a set of parentheses are balanced, 0 else.

   (As implemented in the RASP paper.)

  Example usage:
    shuffle_dyck2 = make_shuffle_dyck(pairs=["()", "{}"])
    shuffle_dyck2("({)}")
    >> [1, 1, 1, 1]
    shuffle_dyck2("(){)}")
    >> [0, 0, 0, 0, 0]

  Args:
    pairs: List of pairs of open and close tokens that each should be balanced.
  """
  assert len(pairs) >= 1

  # Compute running balance of each type of parenthesis
  balances = []
  for pair in pairs:
    assert len(pair) == 2
    open_token, close_token = pair
    balance = make_pair_balance(
        rasp.tokens, open_token=open_token,
        close_token=close_token).named(f"balance_{pair}")
    balances.append(balance)

  # Check if balances where negative anywhere -> parentheses not balanced
  any_negative = balances[0] < 0
  for balance in balances[1:]:
    any_negative = any_negative | (balance < 0)

  # Convert to numerical SOp
  any_negative = rasp.numerical(rasp.Map(lambda x: x,
                                         any_negative)).named("any_negative")

  select_all = rasp.Select(rasp.indices, rasp.indices,
                           rasp.Comparison.TRUE).named("select_all")
  has_neg = rasp.numerical(rasp.Aggregate(select_all, any_negative,
                                          default=0)).named("has_neg")

  # Check if all balances are 0 at the end -> closed all parentheses
  all_zero = balances[0] == 0
  for balance in balances[1:]:
    all_zero = all_zero & (balance == 0)

  select_last = rasp.Select(rasp.indices, length - 1,
                            rasp.Comparison.EQ).named("select_last")
  last_zero = rasp.Aggregate(select_last, all_zero).named("last_zero")

  not_has_neg = (~has_neg).named("not_has_neg")
  return (last_zero & not_has_neg).named("shuffle_dyck")

def make_shuffle_dyck2() -> rasp.SOp:
  return make_shuffle_dyck(pairs=["()", "{}"]).named("shuffle_dyck2")

def make_hist() -> rasp.SOp:
  """Returns the number of times each token occurs in the input.

   (As implemented in the RASP paper.)

  Example usage:
    hist = make_hist()
    hist("abac")
    >> [2, 1, 2, 1]
  """

  same_tok = rasp.Select(rasp.tokens, rasp.tokens,
                         rasp.Comparison.EQ).named("same_tok")
  return rasp.SelectorWidth(same_tok).named("hist")


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP programs only using the subset of RASP supported by the compiler."""

from typing import List, Sequence

from tracr.rasp import rasp

### Programs that work only under non-causal evaluation.

def make_length() -> rasp.SOp:
  """Creates the `length` SOp using selector width primitive.

  Example usage:
    length = make_length()
    length("abcdefg")
    >> [7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0]

  Returns:
    length: SOp mapping an input to a sequence, where every element
      is the length of that sequence.
  """
  all_true_selector = rasp.Select(
      rasp.tokens, rasp.tokens, rasp.Comparison.TRUE).named("all_true_selector")
  return rasp.SelectorWidth(all_true_selector).named("length")

length = make_length()

def make_reverse(sop: rasp.SOp) -> rasp.SOp:
  """Create an SOp that reverses a sequence, using length primitive.

  Example usage:
    reverse = make_reverse(rasp.tokens)
    reverse("Hello")
    >> ['o', 'l', 'l', 'e', 'H']

  Args:
    sop: an SOp

  Returns:
    reverse : SOp that reverses the input sequence.
  """
  opp_idx = (length - rasp.indices).named("opp_idx")
  opp_idx = (opp_idx - 1).named("opp_idx-1")
  reverse_selector = rasp.Select(rasp.indices, opp_idx,
                                 rasp.Comparison.EQ).named("reverse_selector")
  return rasp.Aggregate(reverse_selector, sop).named("reverse")

def make_pair_balance(sop: rasp.SOp, open_token: str,
                      close_token: str) -> rasp.SOp:
  """Return fraction of previous open tokens minus the fraction of close tokens.

   (As implemented in the RASP paper.)

  If the outputs are always non-negative and end in 0, that implies the input
  has balanced parentheses.

  Example usage:
    num_l = make_pair_balance(rasp.tokens, "(", ")")
    num_l("a()b(c))")
    >> [0, 1/2, 0, 0, 1/5, 1/6, 0, -1/8]

  Args:
    sop: Input SOp.
    open_token: Token that counts positive.
    close_token: Token that counts negative.

  Returns:
    pair_balance: SOp mapping an input to a sequence, where every element
      is the fraction of previous open tokens minus previous close tokens.
  """
  bools_open = rasp.numerical(sop == open_token).named("bools_open")
  opens = rasp.numerical(make_frac_prevs(bools_open)).named("opens")

  bools_close = rasp.numerical(sop == close_token).named("bools_close")
  closes = rasp.numerical(make_frac_prevs(bools_close)).named("closes")

  pair_balance = rasp.numerical(rasp.LinearSequenceMap(opens, closes, 1, -1))
  return pair_balance.named("pair_balance")

def make_shuffle_dyck(pairs: List[str]) -> rasp.SOp:
  """Returns 1 if a set of parentheses are balanced, 0 else.

   (As implemented in the RASP paper.)

  Example usage:
    shuffle_dyck2 = make_shuffle_dyck(pairs=["()", "{}"])
    shuffle_dyck2("({)}")
    >> [1, 1, 1, 1]
    shuffle_dyck2("(){)}")
    >> [0, 0, 0, 0, 0]

  Args:
    pairs: List of pairs of open and close tokens that each should be balanced.
  """
  assert len(pairs) >= 1

  # Compute running balance of each type of parenthesis
  balances = []
  for pair in pairs:
    assert len(pair) == 2
    open_token, close_token = pair
    balance = make_pair_balance(
        rasp.tokens, open_token=open_token,
        close_token=close_token).named(f"balance_{pair}")
    balances.append(balance)

  # Check if balances where negative anywhere -> parentheses not balanced
  any_negative = balances[0] < 0
  for balance in balances[1:]:
    any_negative = any_negative | (balance < 0)

  # Convert to numerical SOp
  any_negative = rasp.numerical(rasp.Map(lambda x: x,
                                         any_negative)).named("any_negative")

  select_all = rasp.Select(rasp.indices, rasp.indices,
                           rasp.Comparison.TRUE).named("select_all")
  has_neg = rasp.numerical(rasp.Aggregate(select_all, any_negative,
                                          default=0)).named("has_neg")

  # Check if all balances are 0 at the end -> closed all parentheses
  all_zero = balances[0] == 0
  for balance in balances[1:]:
    all_zero = all_zero & (balance == 0)

  select_last = rasp.Select(rasp.indices, length - 1,
                            rasp.Comparison.EQ).named("select_last")
  last_zero = rasp.Aggregate(select_last, all_zero).named("last_zero")

  not_has_neg = (~has_neg).named("not_has_neg")
  return (last_zero & not_has_neg).named("shuffle_dyck")

def make_shuffle_dyck2() -> rasp.SOp:
  return make_shuffle_dyck(pairs=["()", "{}"]).named("shuffle_dyck2")

def make_hist() -> rasp.SOp:
  """Returns the number of times each token occurs in the input.

   (As implemented in the RASP paper.)

  Example usage:
    hist = make_hist()
    hist("abac")
    >> [2, 1, 2, 1]
  """
  same_tok = rasp.Select(rasp.tokens, rasp.tokens,
                         rasp.Comparison.EQ).named("same_tok")
  return rasp.SelectorWidth(same_tok).named("hist")

def make_sort_unique(vals: rasp.SOp, keys: rasp.SOp) -> rasp.SOp:
  """Returns vals sorted by < relation on keys.

  Only supports unique keys.

  Example usage:
    sort = make_sort(rasp.tokens, rasp.tokens)
    sort([2, 4, 3, 1])
    >> [1, 2, 3, 4]

  Args:
    vals: Values to sort.
    keys: Keys for sorting.
  """

  smaller = rasp.Select(keys, keys, rasp.Comparison.LT).named("smaller")
  target_pos = rasp.SelectorWidth(smaller).named("target_pos")
  sel_new = rasp.Select(target_pos, rasp.indices, rasp.Comparison.EQ)
  return rasp.Aggregate(sel_new, vals).named("sort")


Prompt:
# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RASP programs only using the subset of RASP supported by the compiler."""

from typing import List, Sequence

from tracr.rasp import rasp

### Programs that work only under non-causal evaluation.

def make_length() -> rasp.SOp:
  """Creates the `length` SOp using selector width primitive.

  Example usage:
    length = make_length()
    length("abcdefg")
    >> [7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0]

  Returns:
    length: SOp mapping an input to a sequence, where every element
      is the length of that sequence.
  """
  all_true_selector = rasp.Select(
      rasp.tokens, rasp.tokens, rasp.Comparison.TRUE).named("all_true_selector")
  return rasp.SelectorWidth(all_true_selector).named("length")

length = make_length()

def make_reverse(sop: rasp.SOp) -> rasp.SOp:
  """Create an SOp that reverses a sequence, using length primitive.

  Example usage:
    reverse = make_reverse(rasp.tokens)
    reverse("Hello")
    >> ['o', 'l', 'l', 'e', 'H']

  Args:
    sop: an SOp

  Returns:
    reverse : SOp that reverses the input sequence.
  """
  opp_idx = (length - rasp.indices).named("opp_idx")
  opp_idx = (opp_idx - 1).named("opp_idx-1")
  reverse_selector = rasp.Select(rasp.indices, opp_idx,
                                 rasp.Comparison.EQ).named("reverse_selector")
  return rasp.Aggregate(reverse_selector, sop).named("reverse")

def make_pair_balance(sop: rasp.SOp, open_token: str,
                      close_token: str) -> rasp.SOp:
  """Return fraction of previous open tokens minus the fraction of close tokens.

   (As implemented in the RASP paper.)

  If the outputs are always non-negative and end in 0, that implies the input
  has balanced parentheses.

  Example usage:
    num_l = make_pair_balance(rasp.tokens, "(", ")")
    num_l("a()b(c))")
    >> [0, 1/2, 0, 0, 1/5, 1/6, 0, -1/8]

  Args:
    sop: Input SOp.
    open_token: Token that counts positive.
    close_token: Token that counts negative.

  Returns:
    pair_balance: SOp mapping an input to a sequence, where every element
      is the fraction of previous open tokens minus previous close tokens.
  """
  bools_open = rasp.numerical(sop == open_token).named("bools_open")
  opens = rasp.numerical(make_frac_prevs(bools_open)).named("opens")

  bools_close = rasp.numerical(sop == close_token).named("bools_close")
  closes = rasp.numerical(make_frac_prevs(bools_close)).named("closes")

  pair_balance = rasp.numerical(rasp.LinearSequenceMap(opens, closes, 1, -1))
  return pair_balance.named("pair_balance")

def make_shuffle_dyck(pairs: List[str]) -> rasp.SOp:
  """Returns 1 if a set of parentheses are balanced, 0 else.

   (As implemented in the RASP paper.)

  Example usage:
    shuffle_dyck2 = make_shuffle_dyck(pairs=["()", "{}"])
    shuffle_dyck2("({)}")
    >> [1, 1, 1, 1]
    shuffle_dyck2("(){)}")
    >> [0, 0, 0, 0, 0]

  Args:
    pairs: List of pairs of open and close tokens that each should be balanced.
  """
  assert len(pairs) >= 1

  # Compute running balance of each type of parenthesis
  balances = []
  for pair in pairs:
    assert len(pair) == 2
    open_token, close_token = pair
    balance = make_pair_balance(
        rasp.tokens, open_token=open_token,
        close_token=close_token).named(f"balance_{pair}")
    balances.append(balance)

  # Check if balances where negative anywhere -> parentheses not balanced
  any_negative = balances[0] < 0
  for balance in balances[1:]:
    any_negative = any_negative | (balance < 0)

  # Convert to numerical SOp
  any_negative = rasp.numerical(rasp.Map(lambda x: x,
                                         any_negative)).named("any_negative")

  select_all = rasp.Select(rasp.indices, rasp.indices,
                           rasp.Comparison.TRUE).named("select_all")
  has_neg = rasp.numerical(rasp.Aggregate(select_all, any_negative,
                                          default=0)).named("has_neg")

  # Check if all balances are 0 at the end -> closed all parentheses
  all_zero = balances[0] == 0
  for balance in balances[1:]:
    all_zero = all_zero & (balance == 0)

  select_last = rasp.Select(rasp.indices, length - 1,
                            rasp.Comparison.EQ).named("select_last")
  last_zero = rasp.Aggregate(select_last, all_zero).named("last_zero")

  not_has_neg = (~has_neg).named("not_has_neg")
  return (last_zero & not_has_neg).named("shuffle_dyck")

def make_shuffle_dyck2() -> rasp.SOp:
  return make_shuffle_dyck(pairs=["()", "{}"]).named("shuffle_dyck2")

def make_hist() -> rasp.SOp:
  """Returns the number of times each token occurs in the input.

   (As implemented in the RASP paper.)

  Example usage:
    hist = make_hist()
    hist("abac")
    >> [2, 1, 2, 1]
  """
  same_tok = rasp.Select(rasp.tokens, rasp.tokens,
                         rasp.Comparison.EQ).named("same_tok")
  return rasp.SelectorWidth(same_tok).named("hist")

def make_sort_unique(vals: rasp.SOp, keys: rasp.SOp) -> rasp.SOp:
  """Returns vals sorted by < relation on keys.

  Only supports unique keys.

  Example usage:
    sort = make_sort(rasp.tokens, rasp.tokens)
    sort([2, 4, 3, 1])
    >> [1, 2, 3, 4]

  Args:
    vals: Values to sort.
    keys: Keys for sorting.
  """
  smaller = rasp.Select(keys, keys, rasp.Comparison.LT).named("smaller")
  target_pos = rasp.SelectorWidth(smaller).named("target_pos")
  sel_new = rasp.Select(target_pos, rasp.indices, rasp.Comparison.EQ)
  return rasp.Aggregate(sel_new, vals).named("sort")

def make_sort(vals: rasp.SOp, keys: rasp.SOp, *, max_seq_len: int,
              min_key: float) -> rasp.SOp:
  """Returns vals sorted by < relation on keys, which don't need to be unique.

  The implementation differs from the RASP paper, as it avoids using
  compositions of selectors to break ties. Instead, it uses the arguments
  max_seq_len and min_key to ensure the keys are unique.

  Note that this approach only works for numerical keys.

  Example usage:
    sort = make_sort(rasp.tokens, rasp.tokens, 5, 1)
    sort([2, 4, 3, 1])
    >> [1, 2, 3, 4]
    sort([2, 4, 1, 2])
    >> [1, 2, 2, 4]

  Args:
    vals: Values to sort.
    keys: Keys for sorting.
    max_seq_len: Maximum sequence length (used to ensure keys are unique)
    min_key: Minimum key value (used to ensure keys are unique)

  Returns:
    Output SOp of sort program.
  """

  keys = rasp.SequenceMap(lambda x, i: x + min_key * i / max_seq_len, keys,
                          rasp.indices)
  return make_sort_unique(vals, keys)