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.
# ==============================================================================
"""MLPs to compute arbitrary numerical functions by discretising."""

import dataclasses

from typing import Callable, Iterable, List

from tracr.craft import bases
from tracr.craft import transformers
from tracr.craft import vectorspace_fns
from tracr.utils import errors

@dataclasses.dataclass
class DiscretisingLayerMaterials:
  """Provides components for a hidden layer that discretises the input.

  Attributes:
    action: Function acting on basis directions that defines the computation.
    hidden_space: Vector space of the hidden representation of the layer.
    output_values: Set of output values that correspond to the discretisation.
  """
  action: Callable[[bases.BasisDirection], bases.VectorInBasis]
  hidden_space: bases.VectorSpaceWithBasis
  output_values: List[float]

def _get_discretising_layer(input_value_set: Iterable[float],
                            f: Callable[[float],
                                        float], hidden_name: bases.Name,
                            one_direction: bases.BasisDirection,
                            large_number: float) -> DiscretisingLayerMaterials:
  """Creates a hidden layer that discretises the input of f(x) into a value set.

  The input is split up into a distinct region around each value in
  `input_value_set`:

  elements of value set:  v0   |  v1  |  v2  |  v3  |  v4  | ...
  thresholds:                  t0     t1     t2     t3     t4

  The hidden layer has two activations per threshold:
    hidden_k_1 = ReLU(L * (x - threshold[k]) + 1)
    hidden_k_2 = ReLU(L * (x - threshold[k]))

  Note that hidden_k_1 - hidden_k_2 is:
    1                 if x >= threshold[k] + 1/L
    0                 if x <= threshold[k]
    between 0 and 1   if threshold[k] < x < threshold[k] + 1/L

  So as long as we choose L a big enough number, we have
    hidden_k_1 - hidden_k_2 = 1 if x >= threshold[k].
  i.e. we know in which region the input value is.

  Args:
    input_value_set: Set of discrete input values.
    f: Function to approximate.
    hidden_name: Name for hidden dimensions.
    one_direction: Auxiliary dimension that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.

  Returns:
    DiscretisingLayerMaterials containing all components for the layer.
  """
  output_values, sorted_values = [], []
  for x in sorted(input_value_set):
    res = errors.ignoring_arithmetic_errors(f)(x)
    if res is not None:
      output_values.append(res)
      sorted_values.append(x)

  num_vals = len(sorted_values)
  value_thresholds = [
      (sorted_values[i] + sorted_values[i + 1]) / 2 for i in range(num_vals - 1)
  ]

  hidden_directions = [bases.BasisDirection(f"{hidden_name}start")]
  for k in range(1, num_vals):
    dir0 = bases.BasisDirection(hidden_name, (k, 0))
    dir1 = bases.BasisDirection(hidden_name, (k, 1))
    hidden_directions.extend([dir0, dir1])
  hidden_space = bases.VectorSpaceWithBasis(hidden_directions)

  def action(direction: bases.BasisDirection) -> bases.VectorInBasis:
    # hidden_k_0 = ReLU(L * (x - threshold[k]) + 1)
    # hidden_k_1 = ReLU(L * (x - threshold[k]))
    if direction == one_direction:
      hidden = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(f"{hidden_name}start"))
    else:
      hidden = hidden_space.null_vector()
    for k in range(1, num_vals):
      vec0 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 0)))
      vec1 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 1)))
      if direction == one_direction:
        hidden += (1 - large_number * value_thresholds[k - 1]) * vec0
        hidden -= large_number * value_thresholds[k - 1] * vec1
      else:
        hidden += large_number * vec0 + large_number * vec1
    return hidden

  return DiscretisingLayerMaterials(
      action=action, hidden_space=hidden_space, output_values=output_values)

def map_numerical_mlp(
    f: Callable[[float], float],
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    input_value_set: Iterable[float],
    one_space: bases.VectorSpaceWithBasis,
    large_number: float = 100,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes any function of a single variable f(x).

  This is implemented by discretising the input according to input_value_set
  and defining thresholds that determine which part of the input range will
  is allocated to which value in input_value_set.

  elements of value set:  v0   |  v1  |  v2  |  v3  |  v4  | ...
  thresholds:                  t0     t1     t2     t3     t4

  The MLP computes two hidden activations per threshold:
    hidden_k_0 = ReLU(L * (x - threshold[k]) + 1)
    hidden_k_1 = ReLU(L * (x - threshold[k]))

  Note that hidden_k_1 - hidden_k_2 is:
    1                 if x >= threshold[k] + 1/L
    0                 if x <= threshold[k]
    between 0 and 1   if threshold[k] < x < threshold[k] + 1/L

  So as long as we choose L a big enough number, we have
    hidden_k_0 - hidden_k_1 = 1 if x >= threshold[k].

  The MLP then computes the output as:
    output = f(input[0]) +
      sum((hidden_k_0 - hidden_k_1) * (f(input[k]) - f(input[k-1]))
        for all k=0,1,...)

  This sum will be (by a telescoping sums argument)
    f(input[0])      if x <= threshold[0]
    f(input[k])      if threshold[k-1] < x <= threshold[k] for some other k
    f(input[-1])     if x > threshold[-1]
  which approximates f() up to an accuracy given by input_value_set and L.

  Args:
    f: Function to approximate.
    input_space: 1-d vector space that encodes the input x.
    output_space: 1-d vector space to write the output to.
    input_value_set: Set of values the input can take.
    one_space: Auxiliary 1-d vector space that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.
      Note that too large values of L can lead to numerical issues, particularly
      during inference on GPU/TPU.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(input_space, num_dims=1, name="input_space")
  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  input_space = bases.join_vector_spaces(input_space, one_space)
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  discretising_layer = _get_discretising_layer(
      input_value_set=input_value_set,
      f=f,
      hidden_name=hidden_name,
      one_direction=one_space.basis[0],
      large_number=large_number)
  first_layer = vectorspace_fns.Linear.from_action(
      input_space, discretising_layer.hidden_space, discretising_layer.action)

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    # output = sum(
    #     (hidden_k_0 - hidden_k_1) * (f(input[k]) - f(input[k-1]))
    #   for all k)
    if direction.name == f"{hidden_name}start":
      return discretising_layer.output_values[0] * out_vec
    k, i = direction.value
    # add hidden_k_0 and subtract hidden_k_1
    sign = {0: 1, 1: -1}[i]
    return sign * (discretising_layer.output_values[k] -
                   discretising_layer.output_values[k - 1]) * out_vec

  second_layer = vectorspace_fns.Linear.from_action(
      discretising_layer.hidden_space, output_space, second_layer_action)

  return transformers.MLP(first_layer, second_layer)

def map_numerical_to_categorical_mlp(
    f: Callable[[float], float],
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    input_value_set: Iterable[float],
    one_space: bases.VectorSpaceWithBasis,
    large_number: float = 100,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a numerical to a categorical variable.

  Uses a set of possible output values, and rounds f(x) to the closest value
  in this set to create a categorical output variable.

  The output is discretised the same way as in `map_numerical_mlp`.

  Args:
    f: Function to approximate.
    input_space: 1-d vector space that encodes the input x.
    output_space: n-d vector space to write categorical output to. The output
      directions need to encode the possible output values.
    input_value_set: Set of values the input can take.
    one_space: Auxiliary 1-d space that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(input_space, num_dims=1, name="input_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  input_space = bases.join_vector_spaces(input_space, one_space)

  vec_by_out_val = dict()
  for d in output_space.basis:
    # TODO(b/255937603): Do a similar assert in other places where we expect
    # categorical basis directions to encode values.
    assert d.value is not None, ("output directions need to encode "
                                 "possible output values")
    vec_by_out_val[d.value] = output_space.vector_from_basis_direction(d)

  discretising_layer = _get_discretising_layer(
      input_value_set=input_value_set,
      f=f,
      hidden_name=hidden_name,
      one_direction=one_space.basis[0],
      large_number=large_number)

  assert set(discretising_layer.output_values).issubset(
      set(vec_by_out_val.keys()))

  first_layer = vectorspace_fns.Linear.from_action(
      input_space, discretising_layer.hidden_space, discretising_layer.action)

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    """Computes output value and returns corresponding output direction."""

    if direction.name == f"{hidden_name}start":
      return vec_by_out_val[discretising_layer.output_values[0]]
    else:
      k, i = direction.value
      # add hidden_k_0 and subtract hidden_k_1
      sign = {0: 1, 1: -1}[i]
      out_k = discretising_layer.output_values[k]
      out_k_m_1 = discretising_layer.output_values[k - 1]
      return sign * (vec_by_out_val[out_k] - vec_by_out_val[out_k_m_1])


Prompt:
# 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.
# ==============================================================================
"""MLPs to compute arbitrary numerical functions by discretising."""

import dataclasses

from typing import Callable, Iterable, List

from tracr.craft import bases
from tracr.craft import transformers
from tracr.craft import vectorspace_fns
from tracr.utils import errors

@dataclasses.dataclass
class DiscretisingLayerMaterials:
  """Provides components for a hidden layer that discretises the input.

  Attributes:
    action: Function acting on basis directions that defines the computation.
    hidden_space: Vector space of the hidden representation of the layer.
    output_values: Set of output values that correspond to the discretisation.
  """
  action: Callable[[bases.BasisDirection], bases.VectorInBasis]
  hidden_space: bases.VectorSpaceWithBasis
  output_values: List[float]

def _get_discretising_layer(input_value_set: Iterable[float],
                            f: Callable[[float],
                                        float], hidden_name: bases.Name,
                            one_direction: bases.BasisDirection,
                            large_number: float) -> DiscretisingLayerMaterials:
  """Creates a hidden layer that discretises the input of f(x) into a value set.

  The input is split up into a distinct region around each value in
  `input_value_set`:

  elements of value set:  v0   |  v1  |  v2  |  v3  |  v4  | ...
  thresholds:                  t0     t1     t2     t3     t4

  The hidden layer has two activations per threshold:
    hidden_k_1 = ReLU(L * (x - threshold[k]) + 1)
    hidden_k_2 = ReLU(L * (x - threshold[k]))

  Note that hidden_k_1 - hidden_k_2 is:
    1                 if x >= threshold[k] + 1/L
    0                 if x <= threshold[k]
    between 0 and 1   if threshold[k] < x < threshold[k] + 1/L

  So as long as we choose L a big enough number, we have
    hidden_k_1 - hidden_k_2 = 1 if x >= threshold[k].
  i.e. we know in which region the input value is.

  Args:
    input_value_set: Set of discrete input values.
    f: Function to approximate.
    hidden_name: Name for hidden dimensions.
    one_direction: Auxiliary dimension that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.

  Returns:
    DiscretisingLayerMaterials containing all components for the layer.
  """
  output_values, sorted_values = [], []
  for x in sorted(input_value_set):
    res = errors.ignoring_arithmetic_errors(f)(x)
    if res is not None:
      output_values.append(res)
      sorted_values.append(x)

  num_vals = len(sorted_values)
  value_thresholds = [
      (sorted_values[i] + sorted_values[i + 1]) / 2 for i in range(num_vals - 1)
  ]

  hidden_directions = [bases.BasisDirection(f"{hidden_name}start")]
  for k in range(1, num_vals):
    dir0 = bases.BasisDirection(hidden_name, (k, 0))
    dir1 = bases.BasisDirection(hidden_name, (k, 1))
    hidden_directions.extend([dir0, dir1])
  hidden_space = bases.VectorSpaceWithBasis(hidden_directions)

  def action(direction: bases.BasisDirection) -> bases.VectorInBasis:
    # hidden_k_0 = ReLU(L * (x - threshold[k]) + 1)
    # hidden_k_1 = ReLU(L * (x - threshold[k]))
    if direction == one_direction:
      hidden = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(f"{hidden_name}start"))
    else:
      hidden = hidden_space.null_vector()
    for k in range(1, num_vals):
      vec0 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 0)))
      vec1 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 1)))
      if direction == one_direction:
        hidden += (1 - large_number * value_thresholds[k - 1]) * vec0
        hidden -= large_number * value_thresholds[k - 1] * vec1
      else:
        hidden += large_number * vec0 + large_number * vec1
    return hidden

  return DiscretisingLayerMaterials(
      action=action, hidden_space=hidden_space, output_values=output_values)

def map_numerical_mlp(
    f: Callable[[float], float],
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    input_value_set: Iterable[float],
    one_space: bases.VectorSpaceWithBasis,
    large_number: float = 100,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes any function of a single variable f(x).

  This is implemented by discretising the input according to input_value_set
  and defining thresholds that determine which part of the input range will
  is allocated to which value in input_value_set.

  elements of value set:  v0   |  v1  |  v2  |  v3  |  v4  | ...
  thresholds:                  t0     t1     t2     t3     t4

  The MLP computes two hidden activations per threshold:
    hidden_k_0 = ReLU(L * (x - threshold[k]) + 1)
    hidden_k_1 = ReLU(L * (x - threshold[k]))

  Note that hidden_k_1 - hidden_k_2 is:
    1                 if x >= threshold[k] + 1/L
    0                 if x <= threshold[k]
    between 0 and 1   if threshold[k] < x < threshold[k] + 1/L

  So as long as we choose L a big enough number, we have
    hidden_k_0 - hidden_k_1 = 1 if x >= threshold[k].

  The MLP then computes the output as:
    output = f(input[0]) +
      sum((hidden_k_0 - hidden_k_1) * (f(input[k]) - f(input[k-1]))
        for all k=0,1,...)

  This sum will be (by a telescoping sums argument)
    f(input[0])      if x <= threshold[0]
    f(input[k])      if threshold[k-1] < x <= threshold[k] for some other k
    f(input[-1])     if x > threshold[-1]
  which approximates f() up to an accuracy given by input_value_set and L.

  Args:
    f: Function to approximate.
    input_space: 1-d vector space that encodes the input x.
    output_space: 1-d vector space to write the output to.
    input_value_set: Set of values the input can take.
    one_space: Auxiliary 1-d vector space that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.
      Note that too large values of L can lead to numerical issues, particularly
      during inference on GPU/TPU.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(input_space, num_dims=1, name="input_space")
  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  input_space = bases.join_vector_spaces(input_space, one_space)
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  discretising_layer = _get_discretising_layer(
      input_value_set=input_value_set,
      f=f,
      hidden_name=hidden_name,
      one_direction=one_space.basis[0],
      large_number=large_number)
  first_layer = vectorspace_fns.Linear.from_action(
      input_space, discretising_layer.hidden_space, discretising_layer.action)

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    # output = sum(
    #     (hidden_k_0 - hidden_k_1) * (f(input[k]) - f(input[k-1]))
    #   for all k)
    if direction.name == f"{hidden_name}start":
      return discretising_layer.output_values[0] * out_vec
    k, i = direction.value
    # add hidden_k_0 and subtract hidden_k_1
    sign = {0: 1, 1: -1}[i]
    return sign * (discretising_layer.output_values[k] -
                   discretising_layer.output_values[k - 1]) * out_vec

  second_layer = vectorspace_fns.Linear.from_action(
      discretising_layer.hidden_space, output_space, second_layer_action)

  return transformers.MLP(first_layer, second_layer)

def map_numerical_to_categorical_mlp(
    f: Callable[[float], float],
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    input_value_set: Iterable[float],
    one_space: bases.VectorSpaceWithBasis,
    large_number: float = 100,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a numerical to a categorical variable.

  Uses a set of possible output values, and rounds f(x) to the closest value
  in this set to create a categorical output variable.

  The output is discretised the same way as in `map_numerical_mlp`.

  Args:
    f: Function to approximate.
    input_space: 1-d vector space that encodes the input x.
    output_space: n-d vector space to write categorical output to. The output
      directions need to encode the possible output values.
    input_value_set: Set of values the input can take.
    one_space: Auxiliary 1-d space that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(input_space, num_dims=1, name="input_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  input_space = bases.join_vector_spaces(input_space, one_space)

  vec_by_out_val = dict()
  for d in output_space.basis:
    # TODO(b/255937603): Do a similar assert in other places where we expect
    # categorical basis directions to encode values.
    assert d.value is not None, ("output directions need to encode "
                                 "possible output values")
    vec_by_out_val[d.value] = output_space.vector_from_basis_direction(d)

  discretising_layer = _get_discretising_layer(
      input_value_set=input_value_set,
      f=f,
      hidden_name=hidden_name,
      one_direction=one_space.basis[0],
      large_number=large_number)

  assert set(discretising_layer.output_values).issubset(
      set(vec_by_out_val.keys()))

  first_layer = vectorspace_fns.Linear.from_action(
      input_space, discretising_layer.hidden_space, discretising_layer.action)

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    """Computes output value and returns corresponding output direction."""
    if direction.name == f"{hidden_name}start":
      return vec_by_out_val[discretising_layer.output_values[0]]
    else:
      k, i = direction.value
      # add hidden_k_0 and subtract hidden_k_1
      sign = {0: 1, 1: -1}[i]
      out_k = discretising_layer.output_values[k]
      out_k_m_1 = discretising_layer.output_values[k - 1]
      return sign * (vec_by_out_val[out_k] - vec_by_out_val[out_k_m_1])

  second_layer = vectorspace_fns.Linear.from_action(
      discretising_layer.hidden_space, output_space, second_layer_action)

  return transformers.MLP(first_layer, second_layer)

def linear_sequence_map_numerical_mlp(
    input1_basis_direction: bases.BasisDirection,
    input2_basis_direction: bases.BasisDirection,
    output_basis_direction: bases.BasisDirection,
    input1_factor: float,
    input2_factor: float,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes a linear function f(x, y) = a*x + b*y.

  Args:
    input1_basis_direction: Basis direction that encodes the input x.
    input2_basis_direction: Basis direction that encodes the input y.
    output_basis_direction: Basis direction to write the output to.
    input1_factor: Linear factor a for input x.
    input2_factor: Linear factor a for input y.
    hidden_name: Name for hidden dimensions.
  """
  input_space = bases.VectorSpaceWithBasis(
      [input1_basis_direction, input2_basis_direction])
  output_space = bases.VectorSpaceWithBasis([output_basis_direction])
  out_vec = output_space.vector_from_basis_direction(output_basis_direction)

  hidden_directions = [
      bases.BasisDirection(f"{hidden_name}x", 1),
      bases.BasisDirection(f"{hidden_name}x", -1),
      bases.BasisDirection(f"{hidden_name}y", 1),
      bases.BasisDirection(f"{hidden_name}y", -1)
  ]
  hidden_space = bases.VectorSpaceWithBasis(hidden_directions)
  x_pos_vec, x_neg_vec, y_pos_vec, y_neg_vec = (
      hidden_space.vector_from_basis_direction(d) for d in hidden_directions)

  def first_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:

    output = hidden_space.null_vector()
    if direction == input1_basis_direction:
      output += x_pos_vec - x_neg_vec
    if direction == input2_basis_direction:
      output += y_pos_vec - y_neg_vec
    return output


Prompt:
# ==============================================================================
"""MLPs to compute arbitrary numerical functions by discretising."""

import dataclasses

from typing import Callable, Iterable, List

from tracr.craft import bases
from tracr.craft import transformers
from tracr.craft import vectorspace_fns
from tracr.utils import errors

@dataclasses.dataclass
class DiscretisingLayerMaterials:
  """Provides components for a hidden layer that discretises the input.

  Attributes:
    action: Function acting on basis directions that defines the computation.
    hidden_space: Vector space of the hidden representation of the layer.
    output_values: Set of output values that correspond to the discretisation.
  """
  action: Callable[[bases.BasisDirection], bases.VectorInBasis]
  hidden_space: bases.VectorSpaceWithBasis
  output_values: List[float]

def _get_discretising_layer(input_value_set: Iterable[float],
                            f: Callable[[float],
                                        float], hidden_name: bases.Name,
                            one_direction: bases.BasisDirection,
                            large_number: float) -> DiscretisingLayerMaterials:
  """Creates a hidden layer that discretises the input of f(x) into a value set.

  The input is split up into a distinct region around each value in
  `input_value_set`:

  elements of value set:  v0   |  v1  |  v2  |  v3  |  v4  | ...
  thresholds:                  t0     t1     t2     t3     t4

  The hidden layer has two activations per threshold:
    hidden_k_1 = ReLU(L * (x - threshold[k]) + 1)
    hidden_k_2 = ReLU(L * (x - threshold[k]))

  Note that hidden_k_1 - hidden_k_2 is:
    1                 if x >= threshold[k] + 1/L
    0                 if x <= threshold[k]
    between 0 and 1   if threshold[k] < x < threshold[k] + 1/L

  So as long as we choose L a big enough number, we have
    hidden_k_1 - hidden_k_2 = 1 if x >= threshold[k].
  i.e. we know in which region the input value is.

  Args:
    input_value_set: Set of discrete input values.
    f: Function to approximate.
    hidden_name: Name for hidden dimensions.
    one_direction: Auxiliary dimension that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.

  Returns:
    DiscretisingLayerMaterials containing all components for the layer.
  """
  output_values, sorted_values = [], []
  for x in sorted(input_value_set):
    res = errors.ignoring_arithmetic_errors(f)(x)
    if res is not None:
      output_values.append(res)
      sorted_values.append(x)

  num_vals = len(sorted_values)
  value_thresholds = [
      (sorted_values[i] + sorted_values[i + 1]) / 2 for i in range(num_vals - 1)
  ]

  hidden_directions = [bases.BasisDirection(f"{hidden_name}start")]
  for k in range(1, num_vals):
    dir0 = bases.BasisDirection(hidden_name, (k, 0))
    dir1 = bases.BasisDirection(hidden_name, (k, 1))
    hidden_directions.extend([dir0, dir1])
  hidden_space = bases.VectorSpaceWithBasis(hidden_directions)

  def action(direction: bases.BasisDirection) -> bases.VectorInBasis:
    # hidden_k_0 = ReLU(L * (x - threshold[k]) + 1)
    # hidden_k_1 = ReLU(L * (x - threshold[k]))
    if direction == one_direction:
      hidden = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(f"{hidden_name}start"))
    else:
      hidden = hidden_space.null_vector()
    for k in range(1, num_vals):
      vec0 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 0)))
      vec1 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 1)))
      if direction == one_direction:
        hidden += (1 - large_number * value_thresholds[k - 1]) * vec0
        hidden -= large_number * value_thresholds[k - 1] * vec1
      else:
        hidden += large_number * vec0 + large_number * vec1
    return hidden

  return DiscretisingLayerMaterials(
      action=action, hidden_space=hidden_space, output_values=output_values)

def map_numerical_mlp(
    f: Callable[[float], float],
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    input_value_set: Iterable[float],
    one_space: bases.VectorSpaceWithBasis,
    large_number: float = 100,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes any function of a single variable f(x).

  This is implemented by discretising the input according to input_value_set
  and defining thresholds that determine which part of the input range will
  is allocated to which value in input_value_set.

  elements of value set:  v0   |  v1  |  v2  |  v3  |  v4  | ...
  thresholds:                  t0     t1     t2     t3     t4

  The MLP computes two hidden activations per threshold:
    hidden_k_0 = ReLU(L * (x - threshold[k]) + 1)
    hidden_k_1 = ReLU(L * (x - threshold[k]))

  Note that hidden_k_1 - hidden_k_2 is:
    1                 if x >= threshold[k] + 1/L
    0                 if x <= threshold[k]
    between 0 and 1   if threshold[k] < x < threshold[k] + 1/L

  So as long as we choose L a big enough number, we have
    hidden_k_0 - hidden_k_1 = 1 if x >= threshold[k].

  The MLP then computes the output as:
    output = f(input[0]) +
      sum((hidden_k_0 - hidden_k_1) * (f(input[k]) - f(input[k-1]))
        for all k=0,1,...)

  This sum will be (by a telescoping sums argument)
    f(input[0])      if x <= threshold[0]
    f(input[k])      if threshold[k-1] < x <= threshold[k] for some other k
    f(input[-1])     if x > threshold[-1]
  which approximates f() up to an accuracy given by input_value_set and L.

  Args:
    f: Function to approximate.
    input_space: 1-d vector space that encodes the input x.
    output_space: 1-d vector space to write the output to.
    input_value_set: Set of values the input can take.
    one_space: Auxiliary 1-d vector space that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.
      Note that too large values of L can lead to numerical issues, particularly
      during inference on GPU/TPU.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(input_space, num_dims=1, name="input_space")
  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  input_space = bases.join_vector_spaces(input_space, one_space)
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  discretising_layer = _get_discretising_layer(
      input_value_set=input_value_set,
      f=f,
      hidden_name=hidden_name,
      one_direction=one_space.basis[0],
      large_number=large_number)
  first_layer = vectorspace_fns.Linear.from_action(
      input_space, discretising_layer.hidden_space, discretising_layer.action)

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    # output = sum(
    #     (hidden_k_0 - hidden_k_1) * (f(input[k]) - f(input[k-1]))
    #   for all k)
    if direction.name == f"{hidden_name}start":
      return discretising_layer.output_values[0] * out_vec
    k, i = direction.value
    # add hidden_k_0 and subtract hidden_k_1
    sign = {0: 1, 1: -1}[i]
    return sign * (discretising_layer.output_values[k] -
                   discretising_layer.output_values[k - 1]) * out_vec

  second_layer = vectorspace_fns.Linear.from_action(
      discretising_layer.hidden_space, output_space, second_layer_action)

  return transformers.MLP(first_layer, second_layer)

def map_numerical_to_categorical_mlp(
    f: Callable[[float], float],
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    input_value_set: Iterable[float],
    one_space: bases.VectorSpaceWithBasis,
    large_number: float = 100,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a numerical to a categorical variable.

  Uses a set of possible output values, and rounds f(x) to the closest value
  in this set to create a categorical output variable.

  The output is discretised the same way as in `map_numerical_mlp`.

  Args:
    f: Function to approximate.
    input_space: 1-d vector space that encodes the input x.
    output_space: n-d vector space to write categorical output to. The output
      directions need to encode the possible output values.
    input_value_set: Set of values the input can take.
    one_space: Auxiliary 1-d space that must contain 1 in the input.
    large_number: Large number L that determines accuracy of the computation.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(input_space, num_dims=1, name="input_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  input_space = bases.join_vector_spaces(input_space, one_space)

  vec_by_out_val = dict()
  for d in output_space.basis:
    # TODO(b/255937603): Do a similar assert in other places where we expect
    # categorical basis directions to encode values.
    assert d.value is not None, ("output directions need to encode "
                                 "possible output values")
    vec_by_out_val[d.value] = output_space.vector_from_basis_direction(d)

  discretising_layer = _get_discretising_layer(
      input_value_set=input_value_set,
      f=f,
      hidden_name=hidden_name,
      one_direction=one_space.basis[0],
      large_number=large_number)

  assert set(discretising_layer.output_values).issubset(
      set(vec_by_out_val.keys()))

  first_layer = vectorspace_fns.Linear.from_action(
      input_space, discretising_layer.hidden_space, discretising_layer.action)

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    """Computes output value and returns corresponding output direction."""
    if direction.name == f"{hidden_name}start":
      return vec_by_out_val[discretising_layer.output_values[0]]
    else:
      k, i = direction.value
      # add hidden_k_0 and subtract hidden_k_1
      sign = {0: 1, 1: -1}[i]
      out_k = discretising_layer.output_values[k]
      out_k_m_1 = discretising_layer.output_values[k - 1]
      return sign * (vec_by_out_val[out_k] - vec_by_out_val[out_k_m_1])

  second_layer = vectorspace_fns.Linear.from_action(
      discretising_layer.hidden_space, output_space, second_layer_action)

  return transformers.MLP(first_layer, second_layer)

def linear_sequence_map_numerical_mlp(
    input1_basis_direction: bases.BasisDirection,
    input2_basis_direction: bases.BasisDirection,
    output_basis_direction: bases.BasisDirection,
    input1_factor: float,
    input2_factor: float,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes a linear function f(x, y) = a*x + b*y.

  Args:
    input1_basis_direction: Basis direction that encodes the input x.
    input2_basis_direction: Basis direction that encodes the input y.
    output_basis_direction: Basis direction to write the output to.
    input1_factor: Linear factor a for input x.
    input2_factor: Linear factor a for input y.
    hidden_name: Name for hidden dimensions.
  """
  input_space = bases.VectorSpaceWithBasis(
      [input1_basis_direction, input2_basis_direction])
  output_space = bases.VectorSpaceWithBasis([output_basis_direction])
  out_vec = output_space.vector_from_basis_direction(output_basis_direction)

  hidden_directions = [
      bases.BasisDirection(f"{hidden_name}x", 1),
      bases.BasisDirection(f"{hidden_name}x", -1),
      bases.BasisDirection(f"{hidden_name}y", 1),
      bases.BasisDirection(f"{hidden_name}y", -1)
  ]
  hidden_space = bases.VectorSpaceWithBasis(hidden_directions)
  x_pos_vec, x_neg_vec, y_pos_vec, y_neg_vec = (
      hidden_space.vector_from_basis_direction(d) for d in hidden_directions)

  def first_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    output = hidden_space.null_vector()
    if direction == input1_basis_direction:
      output += x_pos_vec - x_neg_vec
    if direction == input2_basis_direction:
      output += y_pos_vec - y_neg_vec
    return output

  first_layer = vectorspace_fns.Linear.from_action(input_space, hidden_space,
                                                   first_layer_action)

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:

    if direction.name == f"{hidden_name}x":
      return input1_factor * direction.value * out_vec
    if direction.name == f"{hidden_name}y":
      return input2_factor * direction.value * out_vec
    return output_space.null_vector()


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

from absl.testing import absltest
from absl.testing import parameterized
import numpy as np
from tracr.craft import bases
from tracr.craft import tests_common
from tracr.craft.chamber import categorical_attn

class CategoricalAttnTest(tests_common.VectorFnTestCase):

  @parameterized.parameters([
      dict(causal=False, input_seq=[1, 2, 3, 4, 5], result_seq=[3, 3, 3, 3, 3]),
      dict(
          causal=True,
          input_seq=[1, 2, 3, 4, 5],
          result_seq=[1, 1.5, 2, 2.5, 3]),
      dict(causal=False, input_seq=[10], result_seq=[10]),
      dict(causal=True, input_seq=[10], result_seq=[10]),
      dict(causal=False, input_seq=[-1, 0, 1], result_seq=[0, 0, 0]),
      dict(causal=True, input_seq=[-1, 0, 1], result_seq=[-1, -0.5, 0]),
  ])
  def test_categorical_attn_can_implement_select_all(self, causal, input_seq,
                                                     result_seq):
    vocab = range(-20, 20)
    input_space = bases.VectorSpaceWithBasis.from_values("input", vocab)

    output_dir = bases.BasisDirection("output")
    output_space = bases.VectorSpaceWithBasis([output_dir])
    output_vec = output_space.vector_from_basis_direction(output_dir)

    bos_dir = bases.BasisDirection("bos_dimension")
    bos_space = bases.VectorSpaceWithBasis([bos_dir])

    one_dir = bases.BasisDirection("one")
    one_space = bases.VectorSpaceWithBasis([one_dir])

    value_dir = bases.BasisDirection("value")
    value_space = bases.VectorSpaceWithBasis([value_dir])

    input_space = bases.join_vector_spaces(input_space, bos_space, one_space)
    value_space = bases.join_vector_spaces(value_space, bos_space)
    residual_space = bases.join_vector_spaces(input_space, value_space,
                                              output_space)
    one_vec = residual_space.vector_from_basis_direction(one_dir)
    bos_vec = residual_space.vector_from_basis_direction(bos_dir)
    value_vec = residual_space.vector_from_basis_direction(value_dir)

    attn = categorical_attn.categorical_attn(
        key_space=input_space,
        query_space=input_space,
        value_space=value_space,
        output_space=output_space,
        bos_space=bos_space,
        one_space=one_space,
        attn_fn=lambda x, y: True,
        causal=causal)

    test_inputs = [bos_vec + one_vec]
    for x in input_seq:
      test_inputs.append(
          residual_space.vector_from_basis_direction(
              bases.BasisDirection("input", x)) + x * value_vec)
    test_inputs = bases.VectorInBasis.stack(test_inputs)

    # Expect the average of all (previous) tokens
    expected_results = [x * output_vec for x in result_seq]
    expected_results = bases.VectorInBasis.stack(expected_results)

    test_outputs = attn.apply(test_inputs).project(output_space)

    self.assertVectorAllClose(
        tests_common.strip_bos_token(test_outputs), expected_results)

  @parameterized.parameters([
      dict(causal=False, input_seq=[1, 2, 3, 4, 5], default=0),
      dict(causal=True, input_seq=[1, 2, 3, 4, 5], default=1),
      dict(causal=False, input_seq=[10], default=2),
      dict(causal=True, input_seq=[10], default=-3),
      dict(causal=False, input_seq=[-1, 0, 1], default=-2),
      dict(causal=True, input_seq=[-1, 0, 1], default=-1),
  ])
  def test_categorical_attn_can_implement_select_none(self, causal, input_seq,
                                                      default):
    vocab = range(-20, 20)
    input_space = bases.VectorSpaceWithBasis.from_values("input", vocab)

    output_dir = bases.BasisDirection("output")
    output_space = bases.VectorSpaceWithBasis([output_dir])
    default_vec = default * output_space.vector_from_basis_direction(output_dir)

    bos_dir = bases.BasisDirection("bos_dimension")
    bos_space = bases.VectorSpaceWithBasis([bos_dir])

    one_dir = bases.BasisDirection("one")
    one_space = bases.VectorSpaceWithBasis([one_dir])

    value_dir = bases.BasisDirection("value")
    value_space = bases.VectorSpaceWithBasis([value_dir])

    input_space = bases.join_vector_spaces(input_space, bos_space, one_space)
    value_space = bases.join_vector_spaces(value_space, bos_space)
    residual_space = bases.join_vector_spaces(input_space, value_space,
                                              output_space)
    value_vec = residual_space.vector_from_basis_direction(value_dir)
    bos_vec = residual_space.vector_from_basis_direction(bos_dir)
    one_vec = residual_space.vector_from_basis_direction(one_dir)

    attn = categorical_attn.categorical_attn(
        key_space=input_space,
        query_space=input_space,
        value_space=value_space,
        output_space=output_space,
        bos_space=bos_space,
        one_space=one_space,
        attn_fn=lambda x, y: False,
        default_output=default_vec,
        causal=causal,
        always_attend_to_bos=False,
        use_bos_for_default_output=True)

    def make_input(x):

      return (one_vec + x * value_vec +
              residual_space.vector_from_basis_direction(
                  bases.BasisDirection("input", x)))


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

from typing import Callable

import numpy as np

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

_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])

def map_categorical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
  """Returns an MLP that encodes any categorical function of a single variable f(x).

  The hidden layer is the identity and output combines this with a lookup table
    output_k = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: space containing the input x.
    output_space: space containing possible outputs.
    operation: A function operating on basis directions.
  """

  def operation_fn(direction):
    if direction in input_space:
      output_direction = operation(direction)
      if output_direction in output_space:
        return output_space.vector_from_basis_direction(output_direction)
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)


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

from typing import Callable

import numpy as np

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

_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])

def map_categorical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
  """Returns an MLP that encodes any categorical function of a single variable f(x).

  The hidden layer is the identity and output combines this with a lookup table
    output_k = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: space containing the input x.
    output_space: space containing possible outputs.
    operation: A function operating on basis directions.
  """

  def operation_fn(direction):

    if direction in input_space:
      output_direction = operation(direction)
      if output_direction in output_space:
        return output_space.vector_from_basis_direction(output_direction)
    return output_space.null_vector()


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.
# ==============================================================================
"""MLP to compute basic linear functions of one-hot encoded integers."""

from typing import Callable

import numpy as np

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

_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])

def map_categorical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
  """Returns an MLP that encodes any categorical function of a single variable f(x).

  The hidden layer is the identity and output combines this with a lookup table
    output_k = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: space containing the input x.
    output_space: space containing possible outputs.
    operation: A function operating on basis directions.
  """

  def operation_fn(direction):
    if direction in input_space:
      output_direction = operation(direction)
      if output_direction in output_space:
        return output_space.vector_from_basis_direction(output_direction)
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def map_categorical_to_numerical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.Value], float],
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a categorical to a numerical variable.

  The hidden layer is the identity and output combines this with a lookup table
    output = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: Vector space containing the input x.
    output_space: Vector space to write the numerical output to.
    operation: A function operating on basis directions.
  """

  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  def operation_fn(direction):
    if direction in input_space:
      return operation(direction.value) * out_vec
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)


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

from typing import Callable

import numpy as np

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

_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])

def map_categorical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
  """Returns an MLP that encodes any categorical function of a single variable f(x).

  The hidden layer is the identity and output combines this with a lookup table
    output_k = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: space containing the input x.
    output_space: space containing possible outputs.
    operation: A function operating on basis directions.
  """

  def operation_fn(direction):
    if direction in input_space:
      output_direction = operation(direction)
      if output_direction in output_space:
        return output_space.vector_from_basis_direction(output_direction)
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def map_categorical_to_numerical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.Value], float],
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a categorical to a numerical variable.

  The hidden layer is the identity and output combines this with a lookup table
    output = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: Vector space containing the input x.
    output_space: Vector space to write the numerical output to.
    operation: A function operating on basis directions.
  """
  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  def operation_fn(direction):

    if direction in input_space:
      return operation(direction.value) * out_vec
    return output_space.null_vector()


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.
# ==============================================================================
"""MLP to compute basic linear functions of one-hot encoded integers."""

from typing import Callable

import numpy as np

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

_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])

def map_categorical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
  """Returns an MLP that encodes any categorical function of a single variable f(x).

  The hidden layer is the identity and output combines this with a lookup table
    output_k = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: space containing the input x.
    output_space: space containing possible outputs.
    operation: A function operating on basis directions.
  """

  def operation_fn(direction):
    if direction in input_space:
      output_direction = operation(direction)
      if output_direction in output_space:
        return output_space.vector_from_basis_direction(output_direction)
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def map_categorical_to_numerical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.Value], float],
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a categorical to a numerical variable.

  The hidden layer is the identity and output combines this with a lookup table
    output = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: Vector space containing the input x.
    output_space: Vector space to write the numerical output to.
    operation: A function operating on basis directions.
  """
  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  def operation_fn(direction):
    if direction in input_space:
      return operation(direction.value) * out_vec
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def sequence_map_categorical_mlp(
    input1_space: bases.VectorSpaceWithBasis,
    input2_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection, bases.BasisDirection],
                        bases.BasisDirection],
    one_space: bases.VectorSpaceWithBasis = _ONE_SPACE,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes a categorical function of two variables f(x, y).

  The hidden layer of the MLP computes the logical and of all input directions
    hidden_i_j = ReLU(x_i+x_j-1)

  And the output combines this with a lookup table
    output_k = sum(f(i, j)*hidden_i_j for all i,j in input space)

  Args:
    input1_space: Vector space containing the input x.
    input2_space: Vector space containing the input y.
    output_space: Vector space to write outputs to.
    operation: A function operating on basis directions.
    one_space: a reserved 1-d space that always contains a 1.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  if not set(input1_space.basis).isdisjoint(input2_space.basis):
    raise ValueError("Input spaces to a SequenceMap must be disjoint. "
                     "If input spaces are the same, use Map instead!")

  input_space = bases.direct_sum(input1_space, input2_space, one_space)

  def to_hidden(x, y):
    return bases.BasisDirection(hidden_name, (x.name, x.value, y.name, y.value))

  def from_hidden(h):

    x_name, x_value, y_name, y_value = h.value
    x_dir = bases.BasisDirection(x_name, x_value)
    y_dir = bases.BasisDirection(y_name, y_value)
    return x_dir, y_dir


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.
# ==============================================================================
"""MLP to compute basic linear functions of one-hot encoded integers."""

from typing import Callable

import numpy as np

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

_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])

def map_categorical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
  """Returns an MLP that encodes any categorical function of a single variable f(x).

  The hidden layer is the identity and output combines this with a lookup table
    output_k = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: space containing the input x.
    output_space: space containing possible outputs.
    operation: A function operating on basis directions.
  """

  def operation_fn(direction):
    if direction in input_space:
      output_direction = operation(direction)
      if output_direction in output_space:
        return output_space.vector_from_basis_direction(output_direction)
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def map_categorical_to_numerical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.Value], float],
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a categorical to a numerical variable.

  The hidden layer is the identity and output combines this with a lookup table
    output = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: Vector space containing the input x.
    output_space: Vector space to write the numerical output to.
    operation: A function operating on basis directions.
  """
  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  def operation_fn(direction):
    if direction in input_space:
      return operation(direction.value) * out_vec
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def sequence_map_categorical_mlp(
    input1_space: bases.VectorSpaceWithBasis,
    input2_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection, bases.BasisDirection],
                        bases.BasisDirection],
    one_space: bases.VectorSpaceWithBasis = _ONE_SPACE,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes a categorical function of two variables f(x, y).

  The hidden layer of the MLP computes the logical and of all input directions
    hidden_i_j = ReLU(x_i+x_j-1)

  And the output combines this with a lookup table
    output_k = sum(f(i, j)*hidden_i_j for all i,j in input space)

  Args:
    input1_space: Vector space containing the input x.
    input2_space: Vector space containing the input y.
    output_space: Vector space to write outputs to.
    operation: A function operating on basis directions.
    one_space: a reserved 1-d space that always contains a 1.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  if not set(input1_space.basis).isdisjoint(input2_space.basis):
    raise ValueError("Input spaces to a SequenceMap must be disjoint. "
                     "If input spaces are the same, use Map instead!")

  input_space = bases.direct_sum(input1_space, input2_space, one_space)

  def to_hidden(x, y):
    return bases.BasisDirection(hidden_name, (x.name, x.value, y.name, y.value))

  def from_hidden(h):
    x_name, x_value, y_name, y_value = h.value
    x_dir = bases.BasisDirection(x_name, x_value)
    y_dir = bases.BasisDirection(y_name, y_value)
    return x_dir, y_dir

  hidden_dir = []
  for dir1 in input1_space.basis:
    for dir2 in input2_space.basis:
      hidden_dir.append(to_hidden(dir1, dir2))
  hidden_space = bases.VectorSpaceWithBasis(hidden_dir)

  def logical_and(direction):

    if direction in one_space:
      out = bases.VectorInBasis(hidden_space.basis,
                                -np.ones(hidden_space.num_dims))
    elif direction in input1_space:
      dir1 = direction
      out = hidden_space.null_vector()
      for dir2 in input2_space.basis:
        out += hidden_space.vector_from_basis_direction(to_hidden(dir1, dir2))
    else:
      dir2 = direction
      out = hidden_space.null_vector()
      for dir1 in input1_space.basis:
        out += hidden_space.vector_from_basis_direction(to_hidden(dir1, dir2))
    return out


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

from typing import Callable

import numpy as np

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

_ONE_SPACE = bases.VectorSpaceWithBasis.from_names(["one"])

def map_categorical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection], bases.BasisDirection],
) -> transformers.MLP:
  """Returns an MLP that encodes any categorical function of a single variable f(x).

  The hidden layer is the identity and output combines this with a lookup table
    output_k = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: space containing the input x.
    output_space: space containing possible outputs.
    operation: A function operating on basis directions.
  """

  def operation_fn(direction):
    if direction in input_space:
      output_direction = operation(direction)
      if output_direction in output_space:
        return output_space.vector_from_basis_direction(output_direction)
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def map_categorical_to_numerical_mlp(
    input_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.Value], float],
) -> transformers.MLP:
  """Returns an MLP to compute f(x) from a categorical to a numerical variable.

  The hidden layer is the identity and output combines this with a lookup table
    output = sum(f(i)*input_i for all i in input space)

  Args:
    input_space: Vector space containing the input x.
    output_space: Vector space to write the numerical output to.
    operation: A function operating on basis directions.
  """
  bases.ensure_dims(output_space, num_dims=1, name="output_space")
  out_vec = output_space.vector_from_basis_direction(output_space.basis[0])

  def operation_fn(direction):
    if direction in input_space:
      return operation(direction.value) * out_vec
    return output_space.null_vector()

  first_layer = vectorspace_fns.Linear.from_action(input_space, output_space,
                                                   operation_fn)

  second_layer = vectorspace_fns.project(output_space, output_space)

  return transformers.MLP(first_layer, second_layer)

def sequence_map_categorical_mlp(
    input1_space: bases.VectorSpaceWithBasis,
    input2_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    operation: Callable[[bases.BasisDirection, bases.BasisDirection],
                        bases.BasisDirection],
    one_space: bases.VectorSpaceWithBasis = _ONE_SPACE,
    hidden_name: bases.Name = "__hidden__",
) -> transformers.MLP:
  """Returns an MLP that encodes a categorical function of two variables f(x, y).

  The hidden layer of the MLP computes the logical and of all input directions
    hidden_i_j = ReLU(x_i+x_j-1)

  And the output combines this with a lookup table
    output_k = sum(f(i, j)*hidden_i_j for all i,j in input space)

  Args:
    input1_space: Vector space containing the input x.
    input2_space: Vector space containing the input y.
    output_space: Vector space to write outputs to.
    operation: A function operating on basis directions.
    one_space: a reserved 1-d space that always contains a 1.
    hidden_name: Name for hidden dimensions.
  """
  bases.ensure_dims(one_space, num_dims=1, name="one_space")

  if not set(input1_space.basis).isdisjoint(input2_space.basis):
    raise ValueError("Input spaces to a SequenceMap must be disjoint. "
                     "If input spaces are the same, use Map instead!")

  input_space = bases.direct_sum(input1_space, input2_space, one_space)

  def to_hidden(x, y):
    return bases.BasisDirection(hidden_name, (x.name, x.value, y.name, y.value))

  def from_hidden(h):
    x_name, x_value, y_name, y_value = h.value
    x_dir = bases.BasisDirection(x_name, x_value)
    y_dir = bases.BasisDirection(y_name, y_value)
    return x_dir, y_dir

  hidden_dir = []
  for dir1 in input1_space.basis:
    for dir2 in input2_space.basis:
      hidden_dir.append(to_hidden(dir1, dir2))
  hidden_space = bases.VectorSpaceWithBasis(hidden_dir)

  def logical_and(direction):
    if direction in one_space:
      out = bases.VectorInBasis(hidden_space.basis,
                                -np.ones(hidden_space.num_dims))
    elif direction in input1_space:
      dir1 = direction
      out = hidden_space.null_vector()
      for dir2 in input2_space.basis:
        out += hidden_space.vector_from_basis_direction(to_hidden(dir1, dir2))
    else:
      dir2 = direction
      out = hidden_space.null_vector()
      for dir1 in input1_space.basis:
        out += hidden_space.vector_from_basis_direction(to_hidden(dir1, dir2))
    return out

  first_layer = vectorspace_fns.Linear.from_action(input_space, hidden_space,
                                                   logical_and)

  def operation_fn(direction):

    dir1, dir2 = from_hidden(direction)
    output_direction = operation(dir1, dir2)
    if output_direction in output_space:
      return output_space.vector_from_basis_direction(output_direction)
    else:
      return output_space.null_vector()


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


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


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)


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)


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


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)


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


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