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.
# ==============================================================================
"""Pieces for making transformers."""

import abc
import dataclasses
from typing import Iterable, List, Optional, Sequence, Union

import numpy as np

from tracr.craft import bases
from tracr.craft import vectorspace_fns

project = vectorspace_fns.project

def _np_softmax(x, axis=-1):
  x_max = np.max(x, axis=axis, keepdims=True)
  return np.exp(x - x_max) / np.sum(np.exp(x - x_max), axis=axis, keepdims=True)

def _np_relu(x):
  return np.where(x > 0, x, 0)

def relu(x: bases.VectorInBasis) -> bases.VectorInBasis:
  return bases.VectorInBasis(x.basis_directions, _np_relu(x.magnitudes))

class Block(abc.ABC):
  """Transformer block, acting on a sequence of vector space elements.

  Attributes:
    residual_space: Vector space that contains all subspaces the Block interacts
      with. This can be either the full residual space of a model or a subspace.
  """
  residual_space: bases.VectorSpaceWithBasis

  @abc.abstractmethod
  def apply(self, x: bases.VectorInBasis) -> bases.VectorInBasis:
    """Applies self to an input."""

@dataclasses.dataclass
class AttentionHead(Block):
  """A transformer attention head."""
  w_qk: vectorspace_fns.ScalarBilinear
  w_ov: vectorspace_fns.Linear
  residual_space: Optional[bases.VectorSpaceWithBasis] = None
  causal: bool = False

  def __post_init__(self):
    """Infer residual stream and typecheck subspaces."""
    if self.residual_space is None:
      self.residual_space = bases.join_vector_spaces(self.w_qk.left_space,
                                                     self.w_qk.right_space,
                                                     self.w_ov.input_space,
                                                     self.w_ov.output_space)

    assert self.w_qk.left_space.issubspace(self.residual_space)
    assert self.w_qk.right_space.issubspace(self.residual_space)
    assert self.w_ov.input_space.issubspace(self.residual_space)
    assert self.w_ov.output_space.issubspace(self.residual_space)

  def apply(self, x: bases.VectorInBasis) -> bases.VectorInBasis:
    assert x in self.residual_space
    # seq_len x query_space
    queries = x.project(self.w_qk.left_space)
    # seq_len x key_space
    keys = x.project(self.w_qk.right_space)

    attn_matrix = queries.magnitudes @ self.w_qk.matrix @ keys.magnitudes.T

    if self.causal:
      # The 1 gives us the matrix above the diagonal.
      mask = np.triu(np.full_like(attn_matrix, -np.inf), 1)
      attn_matrix = attn_matrix + mask

    attn_weights = _np_softmax(attn_matrix)  # seq_len_from, seq_len_to
    values = self.w_ov_residual(x).magnitudes  # seq_len_to, d_model

    magnitudes = attn_weights @ values  # seq_len_from, d_model
    return bases.VectorInBasis(sorted(self.residual_space.basis), magnitudes)

  def w_ov_residual(self, x: bases.VectorInBasis) -> bases.VectorInBasis:
    """Wov but acting on the residual space."""
    x = project(self.residual_space, self.w_ov.input_space)(x)
    out = self.w_ov(x)
    return project(self.w_ov.output_space, self.residual_space)(out)

  @property
  def num_heads(self) -> int:
    return 1

  def as_multi(self) -> "MultiAttentionHead":
    return MultiAttentionHead([self])

@dataclasses.dataclass
class MultiAttentionHead(Block):
  """Applies attention heads in parallel."""
  sub_blocks: List[Union[AttentionHead, "MultiAttentionHead"]]

  def __post_init__(self):
    spaces = [block.residual_space for block in self.sub_blocks]
    self.residual_space, *others = spaces
    assert all(s == self.residual_space for s in others)

  def apply(self, x: bases.VectorInBasis) -> bases.VectorInBasis:
    # each element is seq_len x embedding
    outs = [block.apply(x) for block in self.sub_blocks]
    return bases.VectorInBasis.sum(outs)  # seq_len x embedding

  @property
  def num_heads(self) -> int:
    return sum(sub_block.num_heads for sub_block in self.sub_blocks)

  def heads(self) -> Iterable[AttentionHead]:
    for sub_block in self.sub_blocks:
      if isinstance(sub_block, AttentionHead):
        yield sub_block
      elif isinstance(sub_block, MultiAttentionHead):
        yield from sub_block.heads()
      else:
        raise NotImplementedError()

  def as_multi(self) -> "MultiAttentionHead":
    return self

@dataclasses.dataclass
class MLP(Block):
  """A transformer MLP block."""
  fst: vectorspace_fns.Linear
  snd: vectorspace_fns.Linear
  residual_space: Optional[bases.VectorSpaceWithBasis] = None

  def __post_init__(self):
    """Typecheck subspaces."""
    if self.residual_space is None:
      self.residual_space = bases.join_vector_spaces(self.fst.input_space,
                                                     self.snd.output_space)

    assert self.fst.output_space == self.snd.input_space
    assert self.fst.input_space.issubspace(self.residual_space)
    assert self.snd.output_space.issubspace(self.residual_space)

  def apply(self, x: bases.VectorInBasis) -> bases.VectorInBasis:
    assert x in self.residual_space

    x = project(self.residual_space, self.fst.input_space)(x)
    hidden = self.fst(x)
    hidden = relu(hidden)
    out = self.snd(hidden)
    return project(self.snd.output_space, self.residual_space)(out)

  @classmethod
  def combine_in_parallel(cls, mlps: Sequence["MLP"]) -> "MLP":
    fst = vectorspace_fns.Linear.combine_in_parallel(
        [block.fst for block in mlps])
    snd = vectorspace_fns.Linear.combine_in_parallel(
        [block.snd for block in mlps])
    return cls(fst=fst, snd=snd, residual_space=None)

# Block that fits into a half-layer, without residual connections.
HalfLayerBlock = Union[MLP, AttentionHead, MultiAttentionHead]

@dataclasses.dataclass
class SeriesWithResiduals(Block):
  """A series of blocks with residual connections."""
  blocks: List[HalfLayerBlock]

  def __post_init__(self):
    spaces = [block.residual_space for block in self.blocks]
    self.residual_space = bases.join_vector_spaces(*spaces)

  def apply(self, x: bases.VectorInBasis) -> bases.VectorInBasis:

    x = x.project(self.residual_space)
    for block in self.blocks:
      x_in = x.project(block.residual_space)
      x_out = block.apply(x_in).project(self.residual_space)
      x = x + x_out
    return x


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.
# ==============================================================================
"""Helpers for handling errors in user-provided functions."""

import functools
import logging
from typing import Any, Callable

def ignoring_arithmetic_errors(fun: Callable[..., Any]) -> Callable[..., Any]:
  """Makes fun return None instead of raising ArithmeticError."""

  @functools.wraps(fun)

  def fun_wrapped(*args):
    try:
      return fun(*args)
    except ArithmeticError:
      logging.warning(
          "Encountered arithmetic error in function: for value %s. "
          "Assuming this input will never occur.", str(args))
      return None

  return fun_wrapped


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.
# ==============================================================================
"""Helpers for handling errors in user-provided functions."""

import functools
import logging
from typing import Any, Callable

def ignoring_arithmetic_errors(fun: Callable[..., Any]) -> Callable[..., Any]:
  """Makes fun return None instead of raising ArithmeticError."""

  @functools.wraps(fun)
  def fun_wrapped(*args):

    try:
      return fun(*args)
    except ArithmeticError:
      logging.warning(
          "Encountered arithmetic error in function: for value %s. "
          "Assuming this input will never occur.", str(args))
      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 transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"

    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")


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 transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):

    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params


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 transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):

      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):

      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):

      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).output

    @hk.transform
    def forward(emb, mask):

      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      return transformer(emb, mask).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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).output

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params_no_mlps = self._zero_mlps(params)

    out_zero_activation = forward_zero.apply(params, next(rng), emb, mask)
    out_no_mlps = forward.apply(params_no_mlps, next(rng), emb, mask)

    self._check_layer_naming(params)
    np.testing.assert_allclose(out_zero_activation, out_no_mlps)
    self.assertFalse(np.allclose(out_zero_activation, 0))

  def test_not_setting_embedding_size_produces_same_output_as_default_model(
      self):
    config = model.TransformerConfig(
        num_heads=2,
        num_layers=2,
        key_size=5,
        mlp_hidden_size=64,
        dropout_rate=0.,
        causal=False,
        layer_norm=False)

    @hk.without_apply_rng
    @hk.transform
    def forward_model(emb, mask):
      return model.Transformer(config)(emb, mask).output

    @hk.without_apply_rng
    @hk.transform
    def forward_superposition(emb, mask):
      return compressed_model.CompressedTransformer(config)(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward_model.init(next(rng), emb, mask)
    params_superposition = {
        k.replace("transformer", "compressed_transformer"): v
        for k, v in params.items()
    }

    out_model = forward_model.apply(params, emb, mask)
    out_superposition = forward_superposition.apply(params_superposition, emb,
                                                    mask)

    self._check_layer_naming(params_superposition)
    np.testing.assert_allclose(out_model, out_superposition)

  @parameterized.parameters(
      dict(embedding_size=2, unembed_at_every_layer=True),
      dict(embedding_size=2, unembed_at_every_layer=False),
      dict(embedding_size=6, unembed_at_every_layer=True),
      dict(embedding_size=6, unembed_at_every_layer=False))
  def test_embbeding_size_produces_correct_shape_of_residuals_and_layer_outputs(
      self, embedding_size, unembed_at_every_layer):

    @hk.transform
    def forward(emb, mask):

      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False))
      return transformer(
          emb,
          mask,
          embedding_size=embedding_size,
          unembed_at_every_layer=unembed_at_every_layer,
      )


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 transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import compressed_model
from tracr.transformer import model

class CompressedTransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "compressed_transformer/layer_{i}/mlp/linear_1"
    # For Attention: "compressed_transformer/layer_{i}/attn/key"
    # For Layer Norm: "compressed_transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "compressed_transformer")
      if len(levels) == 1:
        self.assertEqual(list(params[key].keys()), ["w_emb"])
        continue
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).output

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params_no_mlps = self._zero_mlps(params)

    out_zero_activation = forward_zero.apply(params, next(rng), emb, mask)
    out_no_mlps = forward.apply(params_no_mlps, next(rng), emb, mask)

    self._check_layer_naming(params)
    np.testing.assert_allclose(out_zero_activation, out_no_mlps)
    self.assertFalse(np.allclose(out_zero_activation, 0))

  def test_not_setting_embedding_size_produces_same_output_as_default_model(
      self):
    config = model.TransformerConfig(
        num_heads=2,
        num_layers=2,
        key_size=5,
        mlp_hidden_size=64,
        dropout_rate=0.,
        causal=False,
        layer_norm=False)

    @hk.without_apply_rng
    @hk.transform
    def forward_model(emb, mask):
      return model.Transformer(config)(emb, mask).output

    @hk.without_apply_rng
    @hk.transform
    def forward_superposition(emb, mask):
      return compressed_model.CompressedTransformer(config)(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward_model.init(next(rng), emb, mask)
    params_superposition = {
        k.replace("transformer", "compressed_transformer"): v
        for k, v in params.items()
    }

    out_model = forward_model.apply(params, emb, mask)
    out_superposition = forward_superposition.apply(params_superposition, emb,
                                                    mask)

    self._check_layer_naming(params_superposition)
    np.testing.assert_allclose(out_model, out_superposition)

  @parameterized.parameters(
      dict(embedding_size=2, unembed_at_every_layer=True),
      dict(embedding_size=2, unembed_at_every_layer=False),
      dict(embedding_size=6, unembed_at_every_layer=True),
      dict(embedding_size=6, unembed_at_every_layer=False))
  def test_embbeding_size_produces_correct_shape_of_residuals_and_layer_outputs(
      self, embedding_size, unembed_at_every_layer):

    @hk.transform
    def forward(emb, mask):
      transformer = compressed_model.CompressedTransformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False))
      return transformer(
          emb,
          mask,
          embedding_size=embedding_size,
          unembed_at_every_layer=unembed_at_every_layer,
      )

    seq_len = 4
    model_size = 16

    emb = np.random.random((1, seq_len, model_size))
    mask = np.ones((1, seq_len))
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    activations = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)

    for residual in activations.residuals:
      self.assertEqual(residual.shape, (1, seq_len, embedding_size))

    for layer_output in activations.layer_outputs:
      self.assertEqual(layer_output.shape, (1, seq_len, model_size))

  @parameterized.parameters(
      dict(model_size=2, unembed_at_every_layer=True),
      dict(model_size=2, unembed_at_every_layer=False),
      dict(model_size=6, unembed_at_every_layer=True),
      dict(model_size=6, unembed_at_every_layer=False))
  def test_identity_embedding_produces_same_output_as_standard_model(
      self, model_size, unembed_at_every_layer):

    config = model.TransformerConfig(
        num_heads=2,
        num_layers=2,
        key_size=5,
        mlp_hidden_size=64,
        dropout_rate=0.,
        causal=False,
        layer_norm=False)

    @hk.without_apply_rng
    @hk.transform
    def forward_model(emb, mask):
      return model.Transformer(config)(emb, mask).output

    @hk.without_apply_rng
    @hk.transform
    def forward_superposition(emb, mask):

      return compressed_model.CompressedTransformer(config)(
          emb,
          mask,
          embedding_size=model_size,
          unembed_at_every_layer=unembed_at_every_layer).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.
# ==============================================================================
"""Didactic example of an autoregressive Transformer-based language model.

Glossary of shapes:
- B: Batch size.
- T: Sequence length.
- D: Model embedding size.
- H: Number of attention heads.
- V: Vocabulary size.

Forked from: haiku.examples.transformer.model
"""

import collections
import dataclasses
from typing import Callable, List, Optional

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import attention

# hk.Modules are not always callable: github.com/deepmind/dm-haiku/issues/52
# Ideally, we'd want a type:
# CallableHaikuModule = Intersection[Callable[..., jax.Array], hk.Module]
# But Intersection does not exist (yet): github.com/python/typing/issues/213
CallableHaikuModule = Callable[..., jax.Array]

@chex.dataclass
class TransformerOutput:
  layer_outputs: List[jax.Array]  # [B, T, D]
  residuals: List[jax.Array]  # [B, T, D]
  attn_logits: List[jax.Array]  # [B, H, T, T]
  output: jax.Array  # [B, T, D]
  input_embeddings: jax.Array  # [B, T, D]

@dataclasses.dataclass
class TransformerConfig:
  num_heads: int
  num_layers: int
  key_size: int
  mlp_hidden_size: int
  dropout_rate: float
  activation_function: Callable[[jax.Array], jax.Array] = jax.nn.gelu
  layer_norm: bool = True
  causal: bool = False

@dataclasses.dataclass
class Transformer(hk.Module):
  """A transformer stack."""

  config: TransformerConfig
  name: Optional[str] = None

  def __call__(
      self,
      embeddings: jax.Array,  # [B, T, D]
      mask: jax.Array,  # [B, T]
      *,
      use_dropout: bool = True,
  ) -> TransformerOutput:
    """Transforms input embedding sequences to output embedding sequences."""

    def layer_norm(x: jax.Array) -> jax.Array:
      """Applies a unique LayerNorm to x with default settings."""

      if self.config.layer_norm:
        return hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)(x)
      return x


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.
# ==============================================================================
"""Didactic example of an autoregressive Transformer-based language model.

Glossary of shapes:
- B: Batch size.
- T: Sequence length.
- D: Model embedding size.
- H: Number of attention heads.
- V: Vocabulary size.

Forked from: haiku.examples.transformer.model
"""

import collections
import dataclasses
from typing import Callable, List, Optional

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import attention

# hk.Modules are not always callable: github.com/deepmind/dm-haiku/issues/52
# Ideally, we'd want a type:
# CallableHaikuModule = Intersection[Callable[..., jax.Array], hk.Module]
# But Intersection does not exist (yet): github.com/python/typing/issues/213
CallableHaikuModule = Callable[..., jax.Array]

@chex.dataclass
class TransformerOutput:
  layer_outputs: List[jax.Array]  # [B, T, D]
  residuals: List[jax.Array]  # [B, T, D]
  attn_logits: List[jax.Array]  # [B, H, T, T]
  output: jax.Array  # [B, T, D]
  input_embeddings: jax.Array  # [B, T, D]

@dataclasses.dataclass
class TransformerConfig:
  num_heads: int
  num_layers: int
  key_size: int
  mlp_hidden_size: int
  dropout_rate: float
  activation_function: Callable[[jax.Array], jax.Array] = jax.nn.gelu
  layer_norm: bool = True
  causal: bool = False

@dataclasses.dataclass
class Transformer(hk.Module):
  """A transformer stack."""

  config: TransformerConfig
  name: Optional[str] = None

  def __call__(
      self,
      embeddings: jax.Array,  # [B, T, D]
      mask: jax.Array,  # [B, T]
      *,
      use_dropout: bool = True,
  ) -> TransformerOutput:
    """Transforms input embedding sequences to output embedding sequences."""

    def layer_norm(x: jax.Array) -> jax.Array:
      """Applies a unique LayerNorm to x with default settings."""
      if self.config.layer_norm:
        return hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)(x)
      return x

    initializer = hk.initializers.VarianceScaling(2 / self.config.num_layers)
    dropout_rate = self.config.dropout_rate if use_dropout else 0.
    _, seq_len, model_size = embeddings.shape

    # Compute causal mask for autoregressive sequence modelling.
    mask = mask[:, None, None, :]  # [B, H=1, T'=1, T]
    mask = mask.repeat(seq_len, axis=2)  # [B, H=1, T, T]

    if self.config.causal:
      causal_mask = np.ones((1, 1, seq_len, seq_len))  # [B=1, H=1, T, T]
      causal_mask = np.tril(causal_mask)
      mask = mask * causal_mask  # [B, H=1, T, T]

    # Set up activation collection.
    collected = collections.defaultdict(list)

    def collect(**kwargs):
      for k, v in kwargs.items():
        collected[k].append(v)

    residual = embeddings
    for layer in range(self.config.num_layers):
      with hk.experimental.name_scope(f"layer_{layer}"):
        # First the attention block.
        attn_block = attention.MultiHeadAttention(
            num_heads=self.config.num_heads,
            key_size=self.config.key_size,
            model_size=model_size,
            w_init=initializer,
            name="attn")
        attn_in = layer_norm(residual)
        attn_out = attn_block(attn_in, attn_in, attn_in, mask=mask)
        attn_out, attn_logits = attn_out.out, attn_out.logits
        if dropout_rate > 0:
          attn_out = hk.dropout(hk.next_rng_key(), dropout_rate, attn_out)
        residual = residual + attn_out

        collect(
            residuals=residual, layer_outputs=attn_out, attn_logits=attn_logits)

        # Then the dense block.
        with hk.experimental.name_scope("mlp"):
          dense_block = hk.Sequential([
              hk.Linear(
                  self.config.mlp_hidden_size,
                  w_init=initializer,
                  name="linear_1"),
              self.config.activation_function,
              hk.Linear(model_size, w_init=initializer, name="linear_2"),
          ])
        dense_in = layer_norm(residual)
        dense_out = dense_block(dense_in)
        if dropout_rate > 0:
          dense_out = hk.dropout(hk.next_rng_key(), dropout_rate, dense_out)
        residual = residual + dense_out

        collect(residuals=residual, layer_outputs=dense_out)

    return TransformerOutput(
        residuals=collected["residuals"],
        layer_outputs=collected["layer_outputs"],
        attn_logits=collected["attn_logits"],
        output=layer_norm(residual),
        input_embeddings=embeddings,
    )

@chex.dataclass
class CompiledTransformerModelOutput:
  transformer_output: TransformerOutput
  unembedded_output: jax.Array  # [B, T]

@dataclasses.dataclass
class CompiledTransformerModel(hk.Module):
  """A transformer model with one-hot embeddings."""
  transformer: Transformer
  token_embed: CallableHaikuModule
  position_embed: CallableHaikuModule
  unembed: CallableHaikuModule
  use_unembed_argmax: bool
  pad_token: Optional[int] = None

  def embed(self, tokens: jax.Array) -> jax.Array:

    token_embeddings = self.token_embed(tokens)
    positional_embeddings = self.position_embed(jnp.indices(tokens.shape)[-1])
    return token_embeddings + positional_embeddings  # [B, T, D]


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.
# ==============================================================================
"""Didactic example of an autoregressive Transformer-based language model.

Glossary of shapes:
- B: Batch size.
- T: Sequence length.
- D: Model embedding size.
- H: Number of attention heads.
- V: Vocabulary size.

Forked from: haiku.examples.transformer.model
"""

import collections
import dataclasses
from typing import Callable, List, Optional

import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import attention

# hk.Modules are not always callable: github.com/deepmind/dm-haiku/issues/52
# Ideally, we'd want a type:
# CallableHaikuModule = Intersection[Callable[..., jax.Array], hk.Module]
# But Intersection does not exist (yet): github.com/python/typing/issues/213
CallableHaikuModule = Callable[..., jax.Array]

@chex.dataclass
class TransformerOutput:
  layer_outputs: List[jax.Array]  # [B, T, D]
  residuals: List[jax.Array]  # [B, T, D]
  attn_logits: List[jax.Array]  # [B, H, T, T]
  output: jax.Array  # [B, T, D]
  input_embeddings: jax.Array  # [B, T, D]

@dataclasses.dataclass
class TransformerConfig:
  num_heads: int
  num_layers: int
  key_size: int
  mlp_hidden_size: int
  dropout_rate: float
  activation_function: Callable[[jax.Array], jax.Array] = jax.nn.gelu
  layer_norm: bool = True
  causal: bool = False

@dataclasses.dataclass
class Transformer(hk.Module):
  """A transformer stack."""

  config: TransformerConfig
  name: Optional[str] = None

  def __call__(
      self,
      embeddings: jax.Array,  # [B, T, D]
      mask: jax.Array,  # [B, T]
      *,
      use_dropout: bool = True,
  ) -> TransformerOutput:
    """Transforms input embedding sequences to output embedding sequences."""

    def layer_norm(x: jax.Array) -> jax.Array:
      """Applies a unique LayerNorm to x with default settings."""
      if self.config.layer_norm:
        return hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)(x)
      return x

    initializer = hk.initializers.VarianceScaling(2 / self.config.num_layers)
    dropout_rate = self.config.dropout_rate if use_dropout else 0.
    _, seq_len, model_size = embeddings.shape

    # Compute causal mask for autoregressive sequence modelling.
    mask = mask[:, None, None, :]  # [B, H=1, T'=1, T]
    mask = mask.repeat(seq_len, axis=2)  # [B, H=1, T, T]

    if self.config.causal:
      causal_mask = np.ones((1, 1, seq_len, seq_len))  # [B=1, H=1, T, T]
      causal_mask = np.tril(causal_mask)
      mask = mask * causal_mask  # [B, H=1, T, T]

    # Set up activation collection.
    collected = collections.defaultdict(list)

    def collect(**kwargs):
      for k, v in kwargs.items():
        collected[k].append(v)

    residual = embeddings
    for layer in range(self.config.num_layers):
      with hk.experimental.name_scope(f"layer_{layer}"):
        # First the attention block.
        attn_block = attention.MultiHeadAttention(
            num_heads=self.config.num_heads,
            key_size=self.config.key_size,
            model_size=model_size,
            w_init=initializer,
            name="attn")
        attn_in = layer_norm(residual)
        attn_out = attn_block(attn_in, attn_in, attn_in, mask=mask)
        attn_out, attn_logits = attn_out.out, attn_out.logits
        if dropout_rate > 0:
          attn_out = hk.dropout(hk.next_rng_key(), dropout_rate, attn_out)
        residual = residual + attn_out

        collect(
            residuals=residual, layer_outputs=attn_out, attn_logits=attn_logits)

        # Then the dense block.
        with hk.experimental.name_scope("mlp"):
          dense_block = hk.Sequential([
              hk.Linear(
                  self.config.mlp_hidden_size,
                  w_init=initializer,
                  name="linear_1"),
              self.config.activation_function,
              hk.Linear(model_size, w_init=initializer, name="linear_2"),
          ])
        dense_in = layer_norm(residual)
        dense_out = dense_block(dense_in)
        if dropout_rate > 0:
          dense_out = hk.dropout(hk.next_rng_key(), dropout_rate, dense_out)
        residual = residual + dense_out

        collect(residuals=residual, layer_outputs=dense_out)

    return TransformerOutput(
        residuals=collected["residuals"],
        layer_outputs=collected["layer_outputs"],
        attn_logits=collected["attn_logits"],
        output=layer_norm(residual),
        input_embeddings=embeddings,
    )

@chex.dataclass
class CompiledTransformerModelOutput:
  transformer_output: TransformerOutput
  unembedded_output: jax.Array  # [B, T]

@dataclasses.dataclass
class CompiledTransformerModel(hk.Module):
  """A transformer model with one-hot embeddings."""
  transformer: Transformer
  token_embed: CallableHaikuModule
  position_embed: CallableHaikuModule
  unembed: CallableHaikuModule
  use_unembed_argmax: bool
  pad_token: Optional[int] = None

  def embed(self, tokens: jax.Array) -> jax.Array:
    token_embeddings = self.token_embed(tokens)
    positional_embeddings = self.position_embed(jnp.indices(tokens.shape)[-1])
    return token_embeddings + positional_embeddings  # [B, T, D]

  def __call__(
      self,
      tokens: jax.Array,
      use_dropout: bool = True,
  ) -> CompiledTransformerModelOutput:
    """Embed tokens, pass through model, and unembed output."""

    if self.pad_token is None:
      input_mask = jnp.ones_like(tokens)
    else:
      input_mask = (tokens != self.pad_token)
    input_embeddings = self.embed(tokens)

    transformer_output = self.transformer(
        input_embeddings,
        input_mask,
        use_dropout=use_dropout,
    )
    return CompiledTransformerModelOutput(
        transformer_output=transformer_output,
        unembedded_output=self.unembed(
            transformer_output.output,
            use_unembed_argmax=self.use_unembed_argmax,
        ),
    )


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.
# ==============================================================================
"""Modified transformer to learn a linear compression of the residual stream.

CompressedTransformer adds three arguments compared to Transformer:
- embedding_size: the size of the compressed residual stream.
- unembed_at_every_layer: whether to apply the unembedding before applying
    attention and MLP layers
- return_activations: whether to return all model activations rather than just
    the outputs
"""

import collections
import dataclasses
from typing import Optional

import haiku as hk
import jax
import numpy as np

from tracr.transformer import attention
from tracr.transformer import model

@dataclasses.dataclass
class CompressedTransformer(hk.Module):
  """A transformer stack with linearly compressed residual stream."""

  config: model.TransformerConfig
  name: Optional[str] = None

  def __call__(
      self,
      embeddings: jax.Array,  # [B, T, D]
      mask: jax.Array,  # [B, T]
      *,
      use_dropout: bool = True,
      embedding_size: Optional[int] = None,
      unembed_at_every_layer: bool = False,
  ) -> model.TransformerOutput:  # [B, T, D]
    """Transforms input embedding sequences to output embedding sequences.

    Args:
      embeddings: Input embeddings to pass through the model.
      mask: Boolean mask to restrict the inputs the model uses.
      use_dropout: Turns dropout on/off.
      embedding_size: Dimension to compress the residual stream to.
      unembed_at_every_layer: Whether to unembed the residual stream when
        reading the input for every layer (keeping the layer input sizes) or to
        only unembed before the model output (compressing the layer inputs).

    Returns:
      The outputs of the forward pass through the transformer.
    """

    def layer_norm(x: jax.Array) -> jax.Array:
      """Applies a unique LayerNorm to x with default settings."""

      if self.config.layer_norm:
        return hk.LayerNorm(axis=-1, create_scale=True, create_offset=True)(x)
      return x


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.
# ==============================================================================
"""Basic encoder for inputs with a fixed vocabulary."""

import abc
from typing import Any, List, Optional, Sequence

from tracr.craft import bases

class Encoder(abc.ABC):
  """Encodes a list of tokens into a list of inputs for a transformer model.

  The abstract class does not make assumptions on the input and output types,
  and we have different encoders for different input types.
  """

  @abc.abstractmethod
  def encode(self, inputs: List[Any]) -> List[Any]:
    return list()

  @abc.abstractmethod
  def decode(self, encodings: List[Any]) -> List[Any]:
    return list()

  @property
  def pad_token(self) -> Optional[str]:
    return None

  @property
  def bos_token(self) -> Optional[str]:
    return None

  @property
  def pad_encoding(self) -> Optional[int]:
    return None

  @property
  def bos_encoding(self) -> Optional[int]:
    return None

class NumericalEncoder(Encoder):
  """Encodes numerical variables (simply using the identity mapping)."""

  def encode(self, inputs: List[float]) -> List[float]:
    return inputs

  def decode(self, encodings: List[float]) -> List[float]:
    return encodings

class CategoricalEncoder(Encoder):
  """Encodes categorical variables with a fixed vocabulary."""

  def __init__(
      self,
      basis: Sequence[bases.BasisDirection],
      enforce_bos: bool = False,
      bos_token: Optional[str] = None,
      pad_token: Optional[str] = None,
      max_seq_len: Optional[int] = None,
  ):
    """Initialises. If enforce_bos is set, ensures inputs start with it."""

    if enforce_bos and not bos_token:
      raise ValueError("BOS token must be specified if enforcing BOS.")

    self.encoding_map = {}
    for i, direction in enumerate(basis):
      val = direction.value
      self.encoding_map[val] = i

    if bos_token and bos_token not in self.encoding_map:
      raise ValueError("BOS token missing in encoding.")

    if pad_token and pad_token not in self.encoding_map:
      raise ValueError("PAD token missing in encoding.")

    self.enforce_bos = enforce_bos
    self._bos_token = bos_token
    self._pad_token = pad_token
    self._max_seq_len = max_seq_len


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.
# ==============================================================================
"""Basic encoder for inputs with a fixed vocabulary."""

import abc
from typing import Any, List, Optional, Sequence

from tracr.craft import bases

class Encoder(abc.ABC):
  """Encodes a list of tokens into a list of inputs for a transformer model.

  The abstract class does not make assumptions on the input and output types,
  and we have different encoders for different input types.
  """

  @abc.abstractmethod
  def encode(self, inputs: List[Any]) -> List[Any]:
    return list()

  @abc.abstractmethod
  def decode(self, encodings: List[Any]) -> List[Any]:
    return list()

  @property
  def pad_token(self) -> Optional[str]:
    return None

  @property
  def bos_token(self) -> Optional[str]:
    return None

  @property
  def pad_encoding(self) -> Optional[int]:
    return None

  @property
  def bos_encoding(self) -> Optional[int]:
    return None

class NumericalEncoder(Encoder):
  """Encodes numerical variables (simply using the identity mapping)."""

  def encode(self, inputs: List[float]) -> List[float]:
    return inputs

  def decode(self, encodings: List[float]) -> List[float]:
    return encodings

class CategoricalEncoder(Encoder):
  """Encodes categorical variables with a fixed vocabulary."""

  def __init__(
      self,
      basis: Sequence[bases.BasisDirection],
      enforce_bos: bool = False,
      bos_token: Optional[str] = None,
      pad_token: Optional[str] = None,
      max_seq_len: Optional[int] = None,
  ):
    """Initialises. If enforce_bos is set, ensures inputs start with it."""
    if enforce_bos and not bos_token:
      raise ValueError("BOS token must be specified if enforcing BOS.")

    self.encoding_map = {}
    for i, direction in enumerate(basis):
      val = direction.value
      self.encoding_map[val] = i

    if bos_token and bos_token not in self.encoding_map:
      raise ValueError("BOS token missing in encoding.")

    if pad_token and pad_token not in self.encoding_map:
      raise ValueError("PAD token missing in encoding.")

    self.enforce_bos = enforce_bos
    self._bos_token = bos_token
    self._pad_token = pad_token
    self._max_seq_len = max_seq_len

  def encode(self, inputs: List[bases.Value]) -> List[int]:

    if self.enforce_bos and inputs[0] != self.bos_token:
      raise ValueError("First input token must be BOS token. "
                       f"Should be '{self.bos_token}', but was '{inputs[0]}'.")
    if missing := set(inputs) - set(self.encoding_map.keys()):
      raise ValueError(f"Inputs {missing} not found in encoding ",
                       self.encoding_map.keys())
    if self._max_seq_len is not None and len(inputs) > self._max_seq_len:
      raise ValueError(f"inputs={inputs} are longer than the maximum "
                       f"sequence length {self._max_seq_len}")

    return [self.encoding_map[x] for x in inputs]


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.
# ==============================================================================
"""Basic encoder for inputs with a fixed vocabulary."""

import abc
from typing import Any, List, Optional, Sequence

from tracr.craft import bases

class Encoder(abc.ABC):
  """Encodes a list of tokens into a list of inputs for a transformer model.

  The abstract class does not make assumptions on the input and output types,
  and we have different encoders for different input types.
  """

  @abc.abstractmethod
  def encode(self, inputs: List[Any]) -> List[Any]:
    return list()

  @abc.abstractmethod
  def decode(self, encodings: List[Any]) -> List[Any]:
    return list()

  @property
  def pad_token(self) -> Optional[str]:
    return None

  @property
  def bos_token(self) -> Optional[str]:
    return None

  @property
  def pad_encoding(self) -> Optional[int]:
    return None

  @property
  def bos_encoding(self) -> Optional[int]:
    return None

class NumericalEncoder(Encoder):
  """Encodes numerical variables (simply using the identity mapping)."""

  def encode(self, inputs: List[float]) -> List[float]:
    return inputs

  def decode(self, encodings: List[float]) -> List[float]:
    return encodings

class CategoricalEncoder(Encoder):
  """Encodes categorical variables with a fixed vocabulary."""

  def __init__(
      self,
      basis: Sequence[bases.BasisDirection],
      enforce_bos: bool = False,
      bos_token: Optional[str] = None,
      pad_token: Optional[str] = None,
      max_seq_len: Optional[int] = None,
  ):
    """Initialises. If enforce_bos is set, ensures inputs start with it."""
    if enforce_bos and not bos_token:
      raise ValueError("BOS token must be specified if enforcing BOS.")

    self.encoding_map = {}
    for i, direction in enumerate(basis):
      val = direction.value
      self.encoding_map[val] = i

    if bos_token and bos_token not in self.encoding_map:
      raise ValueError("BOS token missing in encoding.")

    if pad_token and pad_token not in self.encoding_map:
      raise ValueError("PAD token missing in encoding.")

    self.enforce_bos = enforce_bos
    self._bos_token = bos_token
    self._pad_token = pad_token
    self._max_seq_len = max_seq_len

  def encode(self, inputs: List[bases.Value]) -> List[int]:
    if self.enforce_bos and inputs[0] != self.bos_token:
      raise ValueError("First input token must be BOS token. "
                       f"Should be '{self.bos_token}', but was '{inputs[0]}'.")
    if missing := set(inputs) - set(self.encoding_map.keys()):
      raise ValueError(f"Inputs {missing} not found in encoding ",
                       self.encoding_map.keys())
    if self._max_seq_len is not None and len(inputs) > self._max_seq_len:
      raise ValueError(f"inputs={inputs} are longer than the maximum "
                       f"sequence length {self._max_seq_len}")

    return [self.encoding_map[x] for x in inputs]

  def decode(self, encodings: List[int]) -> List[bases.Value]:
    """Recover the tokens that corresponds to `ids`. Inverse of __call__."""

    decoding_map = {val: key for key, val in self.encoding_map.items()}
    if missing := set(encodings) - set(decoding_map.keys()):
      raise ValueError(f"Inputs {missing} not found in decoding map ",
                       decoding_map.keys())
    return [decoding_map[x] for x in encodings]


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 transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"

    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")


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 transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):

    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params


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 transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):

      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):

      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):

      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).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.
# ==============================================================================
"""Tests for transformer.model."""

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.transformer import model

class TransformerTest(parameterized.TestCase):

  def _check_layer_naming(self, params):
    # Modules should be named for example
    # For MLPs: "transformer/layer_{i}/mlp/linear_1"
    # For Attention: "transformer/layer_{i}/attn/key"
    # For Layer Norm: "transformer/layer_{i}/layer_norm"
    for key in params.keys():
      levels = key.split("/")
      self.assertEqual(levels[0], "transformer")
      if levels[1].startswith("layer_norm"):
        continue  # output layer norm
      self.assertStartsWith(levels[1], "layer")
      if levels[2] == "mlp":
        self.assertIn(levels[3], {"linear_1", "linear_2"})
      elif levels[2] == "attn":
        self.assertIn(levels[3], {"key", "query", "value", "linear"})
      else:
        self.assertStartsWith(levels[2], "layer_norm")

  def _zero_mlps(self, params):
    for module in params:
      if "mlp" in module:
        for param in params[module]:
          params[module][param] = jnp.zeros_like(params[module][param])
    return params

  @parameterized.parameters(dict(layer_norm=True), dict(layer_norm=False))
  def test_layer_norm(self, layer_norm):
    # input = [1, 1, 1, 1]
    # If layer norm is used, this should give all-0 output for a freshly
    # initialized model because LN will subtract the mean after each layer.
    # Else we expect non-zero outputs.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=layer_norm))
      return transformer(emb, mask).output

    seq_len = 4
    emb = jnp.ones((1, seq_len, 1))
    mask = jnp.ones((1, seq_len))
    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if layer_norm:
      np.testing.assert_allclose(out, 0)
    else:
      self.assertFalse(np.allclose(out, 0))

  @parameterized.parameters(dict(causal=True), dict(causal=False))
  def test_causal_attention(self, causal):
    # input = [0, random, random, random]
    # mask = [1, 0, 1, 1]
    # For causal attention the second token can only attend to the first one, so
    # it should be the same. For non-causal attention all tokens should change.

    @hk.transform
    def forward(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              layer_norm=False,
              causal=causal))
      return transformer(emb, mask).output

    seq_len = 4
    emb = np.random.random((1, seq_len, 1))
    emb[:, 0, :] = 0
    mask = np.array([[1, 0, 1, 1]])
    emb, mask = jnp.array(emb), jnp.array(mask)

    rng = hk.PRNGSequence(1)
    params = forward.init(next(rng), emb, mask)
    params = self._zero_mlps(params)
    out = forward.apply(params, next(rng), emb, mask)

    self._check_layer_naming(params)
    if causal:
      self.assertEqual(0, out[0, 0, 0])
      self.assertEqual(emb[0, 1, 0], out[0, 1, 0])
    else:
      self.assertNotEqual(0, out[0, 0, 0])
      self.assertNotEqual(emb[0, 1, 0], out[0, 1, 0])
    self.assertNotEqual(emb[0, 2, 0], out[0, 2, 0])
    self.assertNotEqual(emb[0, 3, 0], out[0, 3, 0])

  def test_setting_activation_function_to_zero(self):
    # An activation function that always returns zeros should result in the
    # same model output as setting all MLP weights to zero.

    @hk.transform
    def forward_zero(emb, mask):
      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jnp.zeros_like))
      return transformer(emb, mask).output

    @hk.transform
    def forward(emb, mask):

      transformer = model.Transformer(
          model.TransformerConfig(
              num_heads=2,
              num_layers=2,
              key_size=5,
              mlp_hidden_size=64,
              dropout_rate=0.,
              causal=False,
              layer_norm=False,
              activation_function=jax.nn.gelu))
      return transformer(emb, mask).output