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.
# ==============================================================================
"""Functions on vector spaces."""

import abc
import dataclasses
from typing import Callable, Sequence

import numpy as np

from tracr.craft import bases

VectorSpaceWithBasis = bases.VectorSpaceWithBasis
VectorInBasis = bases.VectorInBasis
BasisDirection = bases.BasisDirection

class VectorFunction(abc.ABC):
  """A function that acts on vectors."""

  input_space: VectorSpaceWithBasis
  output_space: VectorSpaceWithBasis

  @abc.abstractmethod
  def __call__(self, x: VectorInBasis) -> VectorInBasis:
    """Evaluates the function."""

class Linear(VectorFunction):
  """A linear function."""

  def __init__(
      self,
      input_space: VectorSpaceWithBasis,
      output_space: VectorSpaceWithBasis,
      matrix: np.ndarray,
  ):
    """Initialises.

    Args:
      input_space: The input vector space.
      output_space: The output vector space.
      matrix: a [input, output] matrix acting in a (sorted) basis.
    """
    self.input_space = input_space
    self.output_space = output_space
    self.matrix = matrix

  def __post_init__(self) -> None:
    output_size, input_size = self.matrix.shape
    assert input_size == self.input_space.num_dims
    assert output_size == self.output_space.num_dims

  def __call__(self, x: VectorInBasis) -> VectorInBasis:
    if x not in self.input_space:
      raise TypeError(f"x={x} not in self.input_space={self.input_space}.")
    return VectorInBasis(
        basis_directions=sorted(self.output_space.basis),
        magnitudes=x.magnitudes @ self.matrix,
    )

  @classmethod
  def from_action(
      cls,
      input_space: VectorSpaceWithBasis,
      output_space: VectorSpaceWithBasis,
      action: Callable[[BasisDirection], VectorInBasis],
  ) -> "Linear":
    """from_action(i, o)(action) creates a Linear."""

    matrix = np.zeros((input_space.num_dims, output_space.num_dims))
    for i, direction in enumerate(input_space.basis):
      out_vector = action(direction)
      if out_vector not in output_space:
        raise TypeError(f"image of {direction} from input_space={input_space} "
                        f"is not in output_space={output_space}")
      matrix[i, :] = out_vector.magnitudes

    return Linear(input_space, output_space, matrix)

  @classmethod
  def combine_in_parallel(cls, fns: Sequence["Linear"]) -> "Linear":
    """Combines multiple parallel linear functions into a single one."""
    joint_input_space = bases.join_vector_spaces(
        *[fn.input_space for fn in fns])
    joint_output_space = bases.join_vector_spaces(
        *[fn.output_space for fn in fns])

    def action(x: bases.BasisDirection) -> bases.VectorInBasis:
      out = joint_output_space.null_vector()
      for fn in fns:
        if x in fn.input_space:
          x_vec = fn.input_space.vector_from_basis_direction(x)
          out += fn(x_vec).project(joint_output_space)
      return out

    return cls.from_action(joint_input_space, joint_output_space, action)

def project(
    from_space: VectorSpaceWithBasis,
    to_space: VectorSpaceWithBasis,
) -> Linear:
  """Creates a projection."""

  def action(direction: bases.BasisDirection) -> VectorInBasis:
    if direction in to_space:
      return to_space.vector_from_basis_direction(direction)
    else:
      return to_space.null_vector()

  return Linear.from_action(from_space, to_space, action=action)

@dataclasses.dataclass
class ScalarBilinear:
  """A scalar-valued bilinear operator."""
  left_space: VectorSpaceWithBasis
  right_space: VectorSpaceWithBasis
  matrix: np.ndarray

  def __post_init__(self):
    """Ensure matrix acts in sorted bases and typecheck sizes."""
    left_size, right_size = self.matrix.shape
    assert left_size == self.left_space.num_dims
    assert right_size == self.right_space.num_dims

  def __call__(self, x: VectorInBasis, y: VectorInBasis) -> float:
    """Describes the action of the operator on vectors."""

    if x not in self.left_space:
      raise TypeError(f"x={x} not in self.left_space={self.left_space}.")
    if y not in self.right_space:
      raise TypeError(f"y={y} not in self.right_space={self.right_space}.")
    return (x.magnitudes.T @ self.matrix @ y.magnitudes).item()


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.
# ==============================================================================
"""Functions on vector spaces."""

import abc
import dataclasses
from typing import Callable, Sequence

import numpy as np

from tracr.craft import bases

VectorSpaceWithBasis = bases.VectorSpaceWithBasis
VectorInBasis = bases.VectorInBasis
BasisDirection = bases.BasisDirection

class VectorFunction(abc.ABC):
  """A function that acts on vectors."""

  input_space: VectorSpaceWithBasis
  output_space: VectorSpaceWithBasis

  @abc.abstractmethod
  def __call__(self, x: VectorInBasis) -> VectorInBasis:
    """Evaluates the function."""

class Linear(VectorFunction):
  """A linear function."""

  def __init__(
      self,
      input_space: VectorSpaceWithBasis,
      output_space: VectorSpaceWithBasis,
      matrix: np.ndarray,
  ):
    """Initialises.

    Args:
      input_space: The input vector space.
      output_space: The output vector space.
      matrix: a [input, output] matrix acting in a (sorted) basis.
    """
    self.input_space = input_space
    self.output_space = output_space
    self.matrix = matrix

  def __post_init__(self) -> None:
    output_size, input_size = self.matrix.shape
    assert input_size == self.input_space.num_dims
    assert output_size == self.output_space.num_dims

  def __call__(self, x: VectorInBasis) -> VectorInBasis:
    if x not in self.input_space:
      raise TypeError(f"x={x} not in self.input_space={self.input_space}.")
    return VectorInBasis(
        basis_directions=sorted(self.output_space.basis),
        magnitudes=x.magnitudes @ self.matrix,
    )

  @classmethod
  def from_action(
      cls,
      input_space: VectorSpaceWithBasis,
      output_space: VectorSpaceWithBasis,
      action: Callable[[BasisDirection], VectorInBasis],
  ) -> "Linear":
    """from_action(i, o)(action) creates a Linear."""

    matrix = np.zeros((input_space.num_dims, output_space.num_dims))
    for i, direction in enumerate(input_space.basis):
      out_vector = action(direction)
      if out_vector not in output_space:
        raise TypeError(f"image of {direction} from input_space={input_space} "
                        f"is not in output_space={output_space}")
      matrix[i, :] = out_vector.magnitudes

    return Linear(input_space, output_space, matrix)

  @classmethod
  def combine_in_parallel(cls, fns: Sequence["Linear"]) -> "Linear":
    """Combines multiple parallel linear functions into a single one."""
    joint_input_space = bases.join_vector_spaces(
        *[fn.input_space for fn in fns])
    joint_output_space = bases.join_vector_spaces(
        *[fn.output_space for fn in fns])

    def action(x: bases.BasisDirection) -> bases.VectorInBasis:
      out = joint_output_space.null_vector()
      for fn in fns:
        if x in fn.input_space:
          x_vec = fn.input_space.vector_from_basis_direction(x)
          out += fn(x_vec).project(joint_output_space)
      return out

    return cls.from_action(joint_input_space, joint_output_space, action)

def project(
    from_space: VectorSpaceWithBasis,
    to_space: VectorSpaceWithBasis,
) -> Linear:
  """Creates a projection."""

  def action(direction: bases.BasisDirection) -> VectorInBasis:
    if direction in to_space:
      return to_space.vector_from_basis_direction(direction)
    else:
      return to_space.null_vector()

  return Linear.from_action(from_space, to_space, action=action)

@dataclasses.dataclass
class ScalarBilinear:
  """A scalar-valued bilinear operator."""
  left_space: VectorSpaceWithBasis
  right_space: VectorSpaceWithBasis
  matrix: np.ndarray

  def __post_init__(self):
    """Ensure matrix acts in sorted bases and typecheck sizes."""
    left_size, right_size = self.matrix.shape
    assert left_size == self.left_space.num_dims
    assert right_size == self.right_space.num_dims

  def __call__(self, x: VectorInBasis, y: VectorInBasis) -> float:
    """Describes the action of the operator on vectors."""
    if x not in self.left_space:
      raise TypeError(f"x={x} not in self.left_space={self.left_space}.")
    if y not in self.right_space:
      raise TypeError(f"y={y} not in self.right_space={self.right_space}.")
    return (x.magnitudes.T @ self.matrix @ y.magnitudes).item()

  @classmethod
  def from_action(
      cls,
      left_space: VectorSpaceWithBasis,
      right_space: VectorSpaceWithBasis,
      action: Callable[[BasisDirection, BasisDirection], float],
  ) -> "ScalarBilinear":
    """from_action(l, r)(action) creates a ScalarBilinear."""

    matrix = np.zeros((left_space.num_dims, right_space.num_dims))
    for i, left_direction in enumerate(left_space.basis):
      for j, right_direction in enumerate(right_space.basis):
        matrix[i, j] = action(left_direction, right_direction)

    return ScalarBilinear(left_space, right_space, matrix)


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):

    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:

    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""

    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":

    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __radd__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {other} + {self}")
    return self + other

  def __sub__(self, other: "VectorInBasis") -> "VectorInBasis":

    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {self} - {other}")
    magnitudes = self.magnitudes - other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __radd__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {other} + {self}")
    return self + other

  def __sub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {self} - {other}")
    magnitudes = self.magnitudes - other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __rsub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {other} - {self}")
    magnitudes = other.magnitudes - self.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __mul__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes * scalar)

  def __rmul__(self, scalar: float) -> "VectorInBasis":
    return self * scalar

  def __truediv__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes / scalar)

  def __neg__(self) -> "VectorInBasis":
    return (-1) * self

  def __eq__(self, other: "VectorInBasis") -> bool:

    return ((self.basis_directions == other.basis_directions) and
            (self.magnitudes.shape == other.magnitudes.shape) and
            (np.all(self.magnitudes == other.magnitudes)))


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __radd__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {other} + {self}")
    return self + other

  def __sub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {self} - {other}")
    magnitudes = self.magnitudes - other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __rsub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {other} - {self}")
    magnitudes = other.magnitudes - self.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __mul__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes * scalar)

  def __rmul__(self, scalar: float) -> "VectorInBasis":
    return self * scalar

  def __truediv__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes / scalar)

  def __neg__(self) -> "VectorInBasis":
    return (-1) * self

  def __eq__(self, other: "VectorInBasis") -> bool:
    return ((self.basis_directions == other.basis_directions) and
            (self.magnitudes.shape == other.magnitudes.shape) and
            (np.all(self.magnitudes == other.magnitudes)))

  @classmethod
  def sum(cls, vectors: Sequence["VectorInBasis"]) -> "VectorInBasis":
    return cls(vectors[0].basis_directions,
               np.sum([x.magnitudes for x in vectors], axis=0))

  @classmethod
  def stack(cls,
            vectors: Sequence["VectorInBasis"],
            axis: int = 0) -> "VectorInBasis":

    for v in vectors[1:]:
      if v.basis_directions != vectors[0].basis_directions:
        raise TypeError(f"Stacking incompatible bases: {vectors[0]} + {v}")
    return cls(vectors[0].basis_directions,
               np.stack([v.magnitudes for v in vectors], axis=axis))


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __radd__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {other} + {self}")
    return self + other

  def __sub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {self} - {other}")
    magnitudes = self.magnitudes - other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __rsub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {other} - {self}")
    magnitudes = other.magnitudes - self.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __mul__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes * scalar)

  def __rmul__(self, scalar: float) -> "VectorInBasis":
    return self * scalar

  def __truediv__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes / scalar)

  def __neg__(self) -> "VectorInBasis":
    return (-1) * self

  def __eq__(self, other: "VectorInBasis") -> bool:
    return ((self.basis_directions == other.basis_directions) and
            (self.magnitudes.shape == other.magnitudes.shape) and
            (np.all(self.magnitudes == other.magnitudes)))

  @classmethod
  def sum(cls, vectors: Sequence["VectorInBasis"]) -> "VectorInBasis":
    return cls(vectors[0].basis_directions,
               np.sum([x.magnitudes for x in vectors], axis=0))

  @classmethod
  def stack(cls,
            vectors: Sequence["VectorInBasis"],
            axis: int = 0) -> "VectorInBasis":
    for v in vectors[1:]:
      if v.basis_directions != vectors[0].basis_directions:
        raise TypeError(f"Stacking incompatible bases: {vectors[0]} + {v}")
    return cls(vectors[0].basis_directions,
               np.stack([v.magnitudes for v in vectors], axis=axis))

  def project(
      self, basis: Union["VectorSpaceWithBasis", Sequence[BasisDirection]]
  ) -> "VectorInBasis":
    """Projects to the basis."""

    if isinstance(basis, VectorSpaceWithBasis):
      basis = basis.basis
    components = []
    for direction in basis:
      if direction in self.basis_directions:
        components.append(
            self.magnitudes[..., self.basis_directions.index(direction)])
      else:
        components.append(np.zeros_like(self.magnitudes[..., 0]))
    return VectorInBasis(list(basis), np.stack(components, axis=-1))


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __radd__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {other} + {self}")
    return self + other

  def __sub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {self} - {other}")
    magnitudes = self.magnitudes - other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __rsub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {other} - {self}")
    magnitudes = other.magnitudes - self.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __mul__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes * scalar)

  def __rmul__(self, scalar: float) -> "VectorInBasis":
    return self * scalar

  def __truediv__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes / scalar)

  def __neg__(self) -> "VectorInBasis":
    return (-1) * self

  def __eq__(self, other: "VectorInBasis") -> bool:
    return ((self.basis_directions == other.basis_directions) and
            (self.magnitudes.shape == other.magnitudes.shape) and
            (np.all(self.magnitudes == other.magnitudes)))

  @classmethod
  def sum(cls, vectors: Sequence["VectorInBasis"]) -> "VectorInBasis":
    return cls(vectors[0].basis_directions,
               np.sum([x.magnitudes for x in vectors], axis=0))

  @classmethod
  def stack(cls,
            vectors: Sequence["VectorInBasis"],
            axis: int = 0) -> "VectorInBasis":
    for v in vectors[1:]:
      if v.basis_directions != vectors[0].basis_directions:
        raise TypeError(f"Stacking incompatible bases: {vectors[0]} + {v}")
    return cls(vectors[0].basis_directions,
               np.stack([v.magnitudes for v in vectors], axis=axis))

  def project(
      self, basis: Union["VectorSpaceWithBasis", Sequence[BasisDirection]]
  ) -> "VectorInBasis":
    """Projects to the basis."""
    if isinstance(basis, VectorSpaceWithBasis):
      basis = basis.basis
    components = []
    for direction in basis:
      if direction in self.basis_directions:
        components.append(
            self.magnitudes[..., self.basis_directions.index(direction)])
      else:
        components.append(np.zeros_like(self.magnitudes[..., 0]))
    return VectorInBasis(list(basis), np.stack(components, axis=-1))

@dataclasses.dataclass
class VectorSpaceWithBasis:
  """A vector subspace in a given basis."""
  basis: Sequence[BasisDirection]

  def __post_init__(self):
    """Keep basis directions sorted."""
    self.basis = sorted(self.basis)

  @property
  def num_dims(self) -> int:
    return len(self.basis)

  def __contains__(self, item: Union[VectorInBasis, BasisDirection]) -> bool:

    if isinstance(item, BasisDirection):
      return item in self.basis

    return set(self.basis) == set(item.basis_directions)


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __radd__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {other} + {self}")
    return self + other

  def __sub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {self} - {other}")
    magnitudes = self.magnitudes - other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __rsub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {other} - {self}")
    magnitudes = other.magnitudes - self.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __mul__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes * scalar)

  def __rmul__(self, scalar: float) -> "VectorInBasis":
    return self * scalar

  def __truediv__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes / scalar)

  def __neg__(self) -> "VectorInBasis":
    return (-1) * self

  def __eq__(self, other: "VectorInBasis") -> bool:
    return ((self.basis_directions == other.basis_directions) and
            (self.magnitudes.shape == other.magnitudes.shape) and
            (np.all(self.magnitudes == other.magnitudes)))

  @classmethod
  def sum(cls, vectors: Sequence["VectorInBasis"]) -> "VectorInBasis":
    return cls(vectors[0].basis_directions,
               np.sum([x.magnitudes for x in vectors], axis=0))

  @classmethod
  def stack(cls,
            vectors: Sequence["VectorInBasis"],
            axis: int = 0) -> "VectorInBasis":
    for v in vectors[1:]:
      if v.basis_directions != vectors[0].basis_directions:
        raise TypeError(f"Stacking incompatible bases: {vectors[0]} + {v}")
    return cls(vectors[0].basis_directions,
               np.stack([v.magnitudes for v in vectors], axis=axis))

  def project(
      self, basis: Union["VectorSpaceWithBasis", Sequence[BasisDirection]]
  ) -> "VectorInBasis":
    """Projects to the basis."""
    if isinstance(basis, VectorSpaceWithBasis):
      basis = basis.basis
    components = []
    for direction in basis:
      if direction in self.basis_directions:
        components.append(
            self.magnitudes[..., self.basis_directions.index(direction)])
      else:
        components.append(np.zeros_like(self.magnitudes[..., 0]))
    return VectorInBasis(list(basis), np.stack(components, axis=-1))

@dataclasses.dataclass
class VectorSpaceWithBasis:
  """A vector subspace in a given basis."""
  basis: Sequence[BasisDirection]

  def __post_init__(self):
    """Keep basis directions sorted."""
    self.basis = sorted(self.basis)

  @property
  def num_dims(self) -> int:
    return len(self.basis)

  def __contains__(self, item: Union[VectorInBasis, BasisDirection]) -> bool:
    if isinstance(item, BasisDirection):
      return item in self.basis

    return set(self.basis) == set(item.basis_directions)

  def issubspace(self, other: "VectorSpaceWithBasis") -> bool:
    return set(self.basis).issubset(set(other.basis))

  def basis_vectors(self) -> Sequence[VectorInBasis]:
    basis_vector_magnitudes = list(np.eye(self.num_dims))
    return [VectorInBasis(self.basis, m) for m in basis_vector_magnitudes]

  def vector_from_basis_direction(
      self, basis_direction: BasisDirection) -> VectorInBasis:
    i = self.basis.index(basis_direction)
    return VectorInBasis(self.basis, np.eye(self.num_dims)[i])

  def null_vector(self) -> VectorInBasis:
    return VectorInBasis(self.basis, np.zeros(self.num_dims))

  @classmethod
  def from_names(cls, names: Sequence[Name]) -> "VectorSpaceWithBasis":
    """Creates a VectorSpace from a list of names for its basis directions."""
    return cls([BasisDirection(n) for n in names])

  @classmethod
  def from_values(
      cls,
      name: Name,
      values: Iterable[Value],
  ) -> "VectorSpaceWithBasis":
    """Creates a VectorSpace from a list of values for its basis directions."""
    return cls([BasisDirection(name, v) for v in values])

def direct_sum(*vs: VectorSpaceWithBasis) -> VectorSpaceWithBasis:
  """Create a direct sum of the vector spaces.

  Assumes the basis elements of all input vector spaces are
  orthogonal to each other. Maintains the order of the bases.

  Args:
    *vs: the vector spaces to sum.

  Returns:
    the combined vector space.

  Raises:
    Value error in case of overlapping bases.
  """
  # Take the union of all the bases:

  total_basis = sum([v.basis for v in vs], [])

  if len(total_basis) != len(set(total_basis)):
    raise ValueError("Overlapping bases!")

  return VectorSpaceWithBasis(total_basis)


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.
# ==============================================================================
"""Vectors and bases."""

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

import numpy as np

Name = Union[int, str]
Value = Union[int, float, bool, str, tuple]

@dataclasses.dataclass(frozen=True)
class BasisDirection:
  """Represents a basis direction (no magnitude) in a vector space.

  Attributes:
    name: a unique name for this direction.
    value: used to hold a value one-hot-encoded by this direction. e.g.,
      [BasisDirection("vs_1", True), BasisDirection("vs_1", False)] would be
      basis directions of a subspace called "vs_1" which one-hot-encodes the
      values True and False. If provided, considered part of the name for the
      purpose of disambiguating directions.
  """
  name: Name
  value: Optional[Value] = None

  def __str__(self):
    if self.value is None:
      return str(self.name)
    return f"{self.name}:{self.value}"

  def __lt__(self, other: "BasisDirection") -> bool:
    try:
      return (self.name, self.value) < (other.name, other.value)
    except TypeError:
      return str(self) < str(other)

@dataclasses.dataclass
class VectorInBasis:
  """A vector (or array of vectors) in a given basis.

  When magnitudes are 1-d, this is a vector.
  When magnitudes are (n+1)-d, this is an array of vectors,
  where the -1th dimension is the basis dimension.
  """
  basis_directions: Sequence[BasisDirection]
  magnitudes: np.ndarray

  def __post_init__(self):
    """Sort basis directions."""
    if len(self.basis_directions) != self.magnitudes.shape[-1]:
      raise ValueError(
          "Last dimension of magnitudes must be the same as number "
          f"of basis directions. Was {len(self.basis_directions)} "
          f"and {self.magnitudes.shape[-1]}.")

    sort_idx = np.argsort(self.basis_directions)
    self.basis_directions = [self.basis_directions[i] for i in sort_idx]
    self.magnitudes = np.take(self.magnitudes, sort_idx, -1)

  def __add__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {self} + {other}")
    magnitudes = self.magnitudes + other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __radd__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Adding incompatible bases: {other} + {self}")
    return self + other

  def __sub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {self} - {other}")
    magnitudes = self.magnitudes - other.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __rsub__(self, other: "VectorInBasis") -> "VectorInBasis":
    if self.basis_directions != other.basis_directions:
      raise TypeError(f"Subtracting incompatible bases: {other} - {self}")
    magnitudes = other.magnitudes - self.magnitudes
    return VectorInBasis(self.basis_directions, magnitudes)

  def __mul__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes * scalar)

  def __rmul__(self, scalar: float) -> "VectorInBasis":
    return self * scalar

  def __truediv__(self, scalar: float) -> "VectorInBasis":
    return VectorInBasis(self.basis_directions, self.magnitudes / scalar)

  def __neg__(self) -> "VectorInBasis":
    return (-1) * self

  def __eq__(self, other: "VectorInBasis") -> bool:
    return ((self.basis_directions == other.basis_directions) and
            (self.magnitudes.shape == other.magnitudes.shape) and
            (np.all(self.magnitudes == other.magnitudes)))

  @classmethod
  def sum(cls, vectors: Sequence["VectorInBasis"]) -> "VectorInBasis":
    return cls(vectors[0].basis_directions,
               np.sum([x.magnitudes for x in vectors], axis=0))

  @classmethod
  def stack(cls,
            vectors: Sequence["VectorInBasis"],
            axis: int = 0) -> "VectorInBasis":
    for v in vectors[1:]:
      if v.basis_directions != vectors[0].basis_directions:
        raise TypeError(f"Stacking incompatible bases: {vectors[0]} + {v}")
    return cls(vectors[0].basis_directions,
               np.stack([v.magnitudes for v in vectors], axis=axis))

  def project(
      self, basis: Union["VectorSpaceWithBasis", Sequence[BasisDirection]]
  ) -> "VectorInBasis":
    """Projects to the basis."""
    if isinstance(basis, VectorSpaceWithBasis):
      basis = basis.basis
    components = []
    for direction in basis:
      if direction in self.basis_directions:
        components.append(
            self.magnitudes[..., self.basis_directions.index(direction)])
      else:
        components.append(np.zeros_like(self.magnitudes[..., 0]))
    return VectorInBasis(list(basis), np.stack(components, axis=-1))

@dataclasses.dataclass
class VectorSpaceWithBasis:
  """A vector subspace in a given basis."""
  basis: Sequence[BasisDirection]

  def __post_init__(self):
    """Keep basis directions sorted."""
    self.basis = sorted(self.basis)

  @property
  def num_dims(self) -> int:
    return len(self.basis)

  def __contains__(self, item: Union[VectorInBasis, BasisDirection]) -> bool:
    if isinstance(item, BasisDirection):
      return item in self.basis

    return set(self.basis) == set(item.basis_directions)

  def issubspace(self, other: "VectorSpaceWithBasis") -> bool:
    return set(self.basis).issubset(set(other.basis))

  def basis_vectors(self) -> Sequence[VectorInBasis]:
    basis_vector_magnitudes = list(np.eye(self.num_dims))
    return [VectorInBasis(self.basis, m) for m in basis_vector_magnitudes]

  def vector_from_basis_direction(
      self, basis_direction: BasisDirection) -> VectorInBasis:
    i = self.basis.index(basis_direction)
    return VectorInBasis(self.basis, np.eye(self.num_dims)[i])

  def null_vector(self) -> VectorInBasis:
    return VectorInBasis(self.basis, np.zeros(self.num_dims))

  @classmethod
  def from_names(cls, names: Sequence[Name]) -> "VectorSpaceWithBasis":
    """Creates a VectorSpace from a list of names for its basis directions."""
    return cls([BasisDirection(n) for n in names])

  @classmethod
  def from_values(
      cls,
      name: Name,
      values: Iterable[Value],
  ) -> "VectorSpaceWithBasis":
    """Creates a VectorSpace from a list of values for its basis directions."""
    return cls([BasisDirection(name, v) for v in values])

def direct_sum(*vs: VectorSpaceWithBasis) -> VectorSpaceWithBasis:
  """Create a direct sum of the vector spaces.

  Assumes the basis elements of all input vector spaces are
  orthogonal to each other. Maintains the order of the bases.

  Args:
    *vs: the vector spaces to sum.

  Returns:
    the combined vector space.

  Raises:
    Value error in case of overlapping bases.
  """
  # Take the union of all the bases:
  total_basis = sum([v.basis for v in vs], [])

  if len(total_basis) != len(set(total_basis)):
    raise ValueError("Overlapping bases!")

  return VectorSpaceWithBasis(total_basis)

def join_vector_spaces(*vs: VectorSpaceWithBasis) -> VectorSpaceWithBasis:
  """Joins a set of vector spaces allowing them to overlap.

  Assumes the basis elements of all input vector spaces are
  orthogonal to each other. Does not maintain the order of the bases but
  sorts them.

  Args:
    *vs: the vector spaces to sum.

  Returns:
    the combined vector space.
  """
  # Take the union of all the bases:

  total_basis = list(set().union(*[set(v.basis) for v in vs]))
  total_basis = sorted(total_basis)
  return VectorSpaceWithBasis(total_basis)


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)


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)


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)


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)


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()


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)


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)


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)