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

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)


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

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)


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

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)

  def __ror__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other | self."""

    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, other, self)
    return Map(lambda x: x or other, self)


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

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)

  def __ror__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other | self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, other, self)
    return Map(lambda x: x or other, self)

class TokensType(SOp):
  """Primitive SOp returning the original input tokens."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "tokens"

  def __repr__(self):
    return "tokens"

class IndicesType(SOp):
  """Primitive SOp returning the position index at each token."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "indices"

  def __repr__(self):
    return "indices"

class LengthType(SOp):
  """Primitive SOp returning the total length of the input."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "length"

  def __repr__(self):
    return "length"

tokens = TokensType()
indices = IndicesType()
length = LengthType()

class Map(SOp):
  """SOp that evaluates the function elementwise on the input SOp.

  Map(lambda x: x + 1, tokens).eval([1, 2, 3]) == [2, 3, 4]
  """

  def __init__(self, f: Callable[[Value], Value], inner: SOp):

    super().__init__()
    self.f = f
    self.inner = inner

    assert isinstance(self.inner, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

    if isinstance(self.inner, Map):
      # Combine the functions into just one.
      inner_f = self.inner.f
      self.f = lambda t: f(inner_f(t))
      self.inner = self.inner.inner


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

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)

  def __ror__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other | self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, other, self)
    return Map(lambda x: x or other, self)

class TokensType(SOp):
  """Primitive SOp returning the original input tokens."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "tokens"

  def __repr__(self):
    return "tokens"

class IndicesType(SOp):
  """Primitive SOp returning the position index at each token."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "indices"

  def __repr__(self):
    return "indices"

class LengthType(SOp):
  """Primitive SOp returning the total length of the input."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "length"

  def __repr__(self):
    return "length"

tokens = TokensType()
indices = IndicesType()
length = LengthType()

class Map(SOp):
  """SOp that evaluates the function elementwise on the input SOp.

  Map(lambda x: x + 1, tokens).eval([1, 2, 3]) == [2, 3, 4]
  """

  def __init__(self, f: Callable[[Value], Value], inner: SOp):
    super().__init__()
    self.f = f
    self.inner = inner

    assert isinstance(self.inner, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

    if isinstance(self.inner, Map):
      # Combine the functions into just one.
      inner_f = self.inner.f
      self.f = lambda t: f(inner_f(t))
      self.inner = self.inner.inner

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.inner]

class SequenceMap(SOp):
  """SOp that evaluates the function elementwise on the two given SOp's.

  SequenceMap(lambda x, y: x - y, length, tokens).eval([1, 2, 3]) == [2, 1, 0]
  """

  def __init__(self, f: Callable[[Value, Value], Value], fst: SOp, snd: SOp):

    super().__init__()

    if fst == snd:
      logging.warning("Creating a SequenceMap with both inputs being the same "
                      "SOp is discouraged. You should use a Map instead.")

    self.f = f
    self.fst = fst
    self.snd = snd
    assert isinstance(self.fst, SOp)
    assert isinstance(self.snd, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)


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

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)

  def __ror__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other | self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, other, self)
    return Map(lambda x: x or other, self)

class TokensType(SOp):
  """Primitive SOp returning the original input tokens."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "tokens"

  def __repr__(self):
    return "tokens"

class IndicesType(SOp):
  """Primitive SOp returning the position index at each token."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "indices"

  def __repr__(self):
    return "indices"

class LengthType(SOp):
  """Primitive SOp returning the total length of the input."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "length"

  def __repr__(self):
    return "length"

tokens = TokensType()
indices = IndicesType()
length = LengthType()

class Map(SOp):
  """SOp that evaluates the function elementwise on the input SOp.

  Map(lambda x: x + 1, tokens).eval([1, 2, 3]) == [2, 3, 4]
  """

  def __init__(self, f: Callable[[Value], Value], inner: SOp):
    super().__init__()
    self.f = f
    self.inner = inner

    assert isinstance(self.inner, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

    if isinstance(self.inner, Map):
      # Combine the functions into just one.
      inner_f = self.inner.f
      self.f = lambda t: f(inner_f(t))
      self.inner = self.inner.inner

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.inner]

class SequenceMap(SOp):
  """SOp that evaluates the function elementwise on the two given SOp's.

  SequenceMap(lambda x, y: x - y, length, tokens).eval([1, 2, 3]) == [2, 1, 0]
  """

  def __init__(self, f: Callable[[Value, Value], Value], fst: SOp, snd: SOp):
    super().__init__()

    if fst == snd:
      logging.warning("Creating a SequenceMap with both inputs being the same "
                      "SOp is discouraged. You should use a Map instead.")

    self.f = f
    self.fst = fst
    self.snd = snd
    assert isinstance(self.fst, SOp)
    assert isinstance(self.snd, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.fst, self.snd]

class LinearSequenceMap(SequenceMap):
  """SOp that evaluates a linear function elementwise on the two given SOp's."""

  def __init__(self, fst: SOp, snd: SOp, fst_fac: float, snd_fac: float):

    super().__init__(fst=fst, snd=snd, f=lambda x, y: fst_fac * x + snd_fac * y)
    self.fst_fac = fst_fac
    self.snd_fac = snd_fac


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

Every object in the RASP language is a function.

The most important type is S-Op, which is a function List[Value] -> List[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)

  def __ror__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other | self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, other, self)
    return Map(lambda x: x or other, self)

class TokensType(SOp):
  """Primitive SOp returning the original input tokens."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "tokens"

  def __repr__(self):
    return "tokens"

class IndicesType(SOp):
  """Primitive SOp returning the position index at each token."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "indices"

  def __repr__(self):
    return "indices"

class LengthType(SOp):
  """Primitive SOp returning the total length of the input."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "length"

  def __repr__(self):
    return "length"

tokens = TokensType()
indices = IndicesType()
length = LengthType()

class Map(SOp):
  """SOp that evaluates the function elementwise on the input SOp.

  Map(lambda x: x + 1, tokens).eval([1, 2, 3]) == [2, 3, 4]
  """

  def __init__(self, f: Callable[[Value], Value], inner: SOp):
    super().__init__()
    self.f = f
    self.inner = inner

    assert isinstance(self.inner, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

    if isinstance(self.inner, Map):
      # Combine the functions into just one.
      inner_f = self.inner.f
      self.f = lambda t: f(inner_f(t))
      self.inner = self.inner.inner

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.inner]

class SequenceMap(SOp):
  """SOp that evaluates the function elementwise on the two given SOp's.

  SequenceMap(lambda x, y: x - y, length, tokens).eval([1, 2, 3]) == [2, 1, 0]
  """

  def __init__(self, f: Callable[[Value, Value], Value], fst: SOp, snd: SOp):
    super().__init__()

    if fst == snd:
      logging.warning("Creating a SequenceMap with both inputs being the same "
                      "SOp is discouraged. You should use a Map instead.")

    self.f = f
    self.fst = fst
    self.snd = snd
    assert isinstance(self.fst, SOp)
    assert isinstance(self.snd, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.fst, self.snd]

class LinearSequenceMap(SequenceMap):
  """SOp that evaluates a linear function elementwise on the two given SOp's."""

  def __init__(self, fst: SOp, snd: SOp, fst_fac: float, snd_fac: float):
    super().__init__(fst=fst, snd=snd, f=lambda x, y: fst_fac * x + snd_fac * y)
    self.fst_fac = fst_fac
    self.snd_fac = snd_fac

class Full(SOp):
  """A SOp evaluating to [fill]*len(input_values)."""

  def __init__(self, fill: Value):
    super().__init__()
    self.fill = fill

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

def sop_not(sop: SOp) -> SOp:
  return Map(lambda t: not t, sop)

class ConstantSOp(SOp, Generic[VT]):
  """A constant S-Op for testing purposes."""

  def __init__(self, value: Sequence[VT], check_length: bool = True):

    super().__init__()
    self.value = value
    self.check_length = check_length


Prompt:
[Value].

An S-Op represents a state inside the residual stream of the transformer.
Therefore, any RASP program that represents a transformer computation must
define a final S-Op that represents the state of the residual stream at the
end of the computation. In particular, given an S-Op `x`,
`x([1, 2, 3])` represents something like the state of the residual stream
at location `x` when the transformer is fed [1, 2, 3] as input.

A secondary (but still important) type is Selector, which is a function
List[Value] -> List[List[bool]]. Given a Selector `sel`, sel([1, 2, 3])
represents something like an attention matrix in the transformer.

For a full reference on RASP, see https://arxiv.org/abs/2106.06981.
"""

import abc
import collections.abc
import copy
import enum
import functools
import itertools
from typing import (Any, Callable, Dict, Generic, List, Mapping, Optional,
                    Sequence, TypeVar, Union)

from absl import logging
import numpy as np
from typing_extensions import Protocol

SelectorValue = List[List[bool]]
NumericValue = Union[int, float]
Value = Union[None, int, float, str, bool]
VT = TypeVar("VT", bound=Value)
RASPExprT = TypeVar("RASPExprT", bound="RASPExpr")
SOpT = TypeVar("SOpT", bound="SOp")
T = TypeVar("T")

_NAME_KEY = "name"
_ENCODING_KEY = "encoding"

# These are run on every expression when it's initialised.
# Add your own annotators to this dict to add custom default annotations.
#
# For example, DEFAULT_ANNOTATORS['foo'] will provide the default value for
# expr.annotations['foo]. The annotator will get called lazily the first time
# that key is accessed.
#
# See the `default_name` annotator for a full example.
DEFAULT_ANNOTATORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)

  def __ror__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other | self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, other, self)
    return Map(lambda x: x or other, self)

class TokensType(SOp):
  """Primitive SOp returning the original input tokens."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "tokens"

  def __repr__(self):
    return "tokens"

class IndicesType(SOp):
  """Primitive SOp returning the position index at each token."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "indices"

  def __repr__(self):
    return "indices"

class LengthType(SOp):
  """Primitive SOp returning the total length of the input."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "length"

  def __repr__(self):
    return "length"

tokens = TokensType()
indices = IndicesType()
length = LengthType()

class Map(SOp):
  """SOp that evaluates the function elementwise on the input SOp.

  Map(lambda x: x + 1, tokens).eval([1, 2, 3]) == [2, 3, 4]
  """

  def __init__(self, f: Callable[[Value], Value], inner: SOp):
    super().__init__()
    self.f = f
    self.inner = inner

    assert isinstance(self.inner, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

    if isinstance(self.inner, Map):
      # Combine the functions into just one.
      inner_f = self.inner.f
      self.f = lambda t: f(inner_f(t))
      self.inner = self.inner.inner

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.inner]

class SequenceMap(SOp):
  """SOp that evaluates the function elementwise on the two given SOp's.

  SequenceMap(lambda x, y: x - y, length, tokens).eval([1, 2, 3]) == [2, 1, 0]
  """

  def __init__(self, f: Callable[[Value, Value], Value], fst: SOp, snd: SOp):
    super().__init__()

    if fst == snd:
      logging.warning("Creating a SequenceMap with both inputs being the same "
                      "SOp is discouraged. You should use a Map instead.")

    self.f = f
    self.fst = fst
    self.snd = snd
    assert isinstance(self.fst, SOp)
    assert isinstance(self.snd, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.fst, self.snd]

class LinearSequenceMap(SequenceMap):
  """SOp that evaluates a linear function elementwise on the two given SOp's."""

  def __init__(self, fst: SOp, snd: SOp, fst_fac: float, snd_fac: float):
    super().__init__(fst=fst, snd=snd, f=lambda x, y: fst_fac * x + snd_fac * y)
    self.fst_fac = fst_fac
    self.snd_fac = snd_fac

class Full(SOp):
  """A SOp evaluating to [fill]*len(input_values)."""

  def __init__(self, fill: Value):
    super().__init__()
    self.fill = fill

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

def sop_not(sop: SOp) -> SOp:
  return Map(lambda t: not t, sop)

class ConstantSOp(SOp, Generic[VT]):
  """A constant S-Op for testing purposes."""

  def __init__(self, value: Sequence[VT], check_length: bool = True):
    super().__init__()
    self.value = value
    self.check_length = check_length

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

### Selectors.

class Predicate(Protocol):

  def __call__(self, key: Value, query: Value) -> bool:
    """Applies the predicate."""

class Comparison(enum.Enum):
  """A two-place boolean comparison predicate for use in Select."""
  EQ = "=="
  LT = "<"
  LEQ = "<="
  GT = ">"
  GEQ = ">="
  NEQ = "!="
  TRUE = "True"
  FALSE = "False"

  def __call__(self, key: Value, query: Value) -> bool:

    if key is None:
      raise ValueError("key is None!")
    if query is None:
      raise ValueError("query is None!")
    return _comparison_table[self](key, query)


Prompt:
ORS: Dict[str, "Annotator"] = {}

class Annotator(Protocol):

  def __call__(self, expr: "RASPExpr") -> Any:
    """What annotation to add to `expr`."""

class _Annotations(collections.abc.Mapping):
  """Holds the expression's annotations.

  It's immutable to the user, but will attempt to generate default values
  lazily when missing keys are requested.
  """

  def __init__(self, expr, **kwargs: Any):
    self._expr = expr
    self._inner_dict: Dict[str, Any] = {**kwargs}

  def __getitem__(self, key: str) -> Any:
    if key not in self._inner_dict:
      if key not in DEFAULT_ANNOTATORS:
        raise KeyError(
            f"No annotation exists for key '{key}'. "
            f"Available keys: {list(*self.keys(), *DEFAULT_ANNOTATORS.keys())}")
      self._inner_dict[key] = DEFAULT_ANNOTATORS[key](self._expr)

    return self._inner_dict[key]

  def __iter__(self):
    return iter(self._inner_dict)

  def __len__(self):
    return len(self._inner_dict)

class RASPExpr(abc.ABC):
  """A class distinguishing RASP expressions from other objects."""
  _ids = itertools.count(1)

  def __init__(self):
    self._annotations: Mapping[str, Any] = _Annotations(self)

  @abc.abstractmethod
  def __call__(self,
               xs: Sequence[Value]) -> Union[Sequence[Value], SelectorValue]:
    """Evaluates the RASPExpr using the standard evaluator."""

  @property
  def annotations(self) -> Mapping[str, Any]:
    """The annotations of this expression instance."""
    return self._annotations

  @annotations.setter
  def annotations(self, annotations: Mapping[str, Any]):
    self._annotations = _Annotations(self, **annotations)

  @property
  def name(self) -> str:
    """The name of this expression."""
    return self.annotations[_NAME_KEY]

  @property
  @abc.abstractmethod
  def children(self) -> Sequence["RASPExpr"]:
    """Direct dependencies of this expression."""

  @functools.cached_property
  def unique_id(self):
    """A unique id for every expression instance."""
    return next(self._ids)

  def copy(self: RASPExprT) -> RASPExprT:
    """Returns a shallow copy of this RASPExpr with a new ID."""
    return copy.copy(self)

  @property
  def label(self) -> str:
    return f"{self.name}_{self.unique_id}"

  def named(self: RASPExprT, name: str) -> RASPExprT:
    """Convenience method for adding a name."""
    return annotate(self, name=name)

  def annotated(self: RASPExprT, **annotations) -> RASPExprT:
    """Convenience method for adding annotations."""
    return annotate(self, **annotations)

def annotate(expr: RASPExprT, **annotations) -> RASPExprT:
  """Creates a new expr with added annotations."""
  new = expr.copy()
  # Note that new annotations will overwrite existing ones with matching keys.
  new.annotations = {**expr.annotations, **annotations}
  return new

### S-Ops.

class SOp(RASPExpr):
  """A Sequence Operation."""

  def __call__(self, xs: Sequence[Value]) -> Sequence[Value]:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of SOps using numeric operators with constant values.
  # Note: if inheriting SOp by a dataclass, make sure to disable eq and order,
  # as they will override these.

  def __lt__(self, other: Value) -> "SOp":
    """self < other."""
    return Map(lambda x: x < other, self)

  def __le__(self, other: Value) -> "SOp":
    """self <= other."""
    return Map(lambda x: x <= other, self)

  def __eq__(self, other: Value) -> "SOp":
    """self == other."""
    return Map(lambda x: x == other, self)

  def __ne__(self, other: Value) -> "SOp":
    """self != other."""
    return Map(lambda x: x != other, self)

  def __gt__(self, other: Value) -> "SOp":
    """self > other."""
    return Map(lambda x: x > other, self)

  def __ge__(self, other: Value) -> "SOp":
    """self >= other."""
    return Map(lambda x: x >= other, self)

  def __add__(self, other: Union["SOp", Value]) -> "SOp":
    """self + other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, self, other)
    return Map(lambda x: x + other, self)

  def __radd__(self, other: Union["SOp", Value]) -> "SOp":
    """other + self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x + y, other, self)
    return Map(lambda x: other + x, self)

  def __sub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self - other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, self, other)
    return Map(lambda x: x - other, self)

  def __rsub__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other - self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x - y, other, self)
    return Map(lambda x: other - x, self)

  def __mul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self * other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, self, other)
    return Map(lambda x: x * other, self)

  def __rmul__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other * self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x * y, other, self)
    return Map(lambda x: other * x, self)

  def __truediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self / other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, self, other)
    return Map(lambda x: x / other, self)

  def __rtruediv__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other / self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x / y, other, self)
    return Map(lambda x: other / x, self)

  def __invert__(self) -> "SOp":
    return Map(lambda x: not x, self)

  def __and__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self & other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, self, other)
    return Map(lambda x: x and other, self)

  def __or__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """self | other."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, self, other)
    return Map(lambda x: x or other, self)

  def __rand__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other & self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x and y, other, self)
    return Map(lambda x: other and x, self)

  def __ror__(self, other: Union["SOp", NumericValue]) -> "SOp":
    """other | self."""
    if isinstance(other, SOp):
      return SequenceMap(lambda x, y: x or y, other, self)
    return Map(lambda x: x or other, self)

class TokensType(SOp):
  """Primitive SOp returning the original input tokens."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "tokens"

  def __repr__(self):
    return "tokens"

class IndicesType(SOp):
  """Primitive SOp returning the position index at each token."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "indices"

  def __repr__(self):
    return "indices"

class LengthType(SOp):
  """Primitive SOp returning the total length of the input."""

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

  @property
  def label(self) -> str:
    return "length"

  def __repr__(self):
    return "length"

tokens = TokensType()
indices = IndicesType()
length = LengthType()

class Map(SOp):
  """SOp that evaluates the function elementwise on the input SOp.

  Map(lambda x: x + 1, tokens).eval([1, 2, 3]) == [2, 3, 4]
  """

  def __init__(self, f: Callable[[Value], Value], inner: SOp):
    super().__init__()
    self.f = f
    self.inner = inner

    assert isinstance(self.inner, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

    if isinstance(self.inner, Map):
      # Combine the functions into just one.
      inner_f = self.inner.f
      self.f = lambda t: f(inner_f(t))
      self.inner = self.inner.inner

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.inner]

class SequenceMap(SOp):
  """SOp that evaluates the function elementwise on the two given SOp's.

  SequenceMap(lambda x, y: x - y, length, tokens).eval([1, 2, 3]) == [2, 1, 0]
  """

  def __init__(self, f: Callable[[Value, Value], Value], fst: SOp, snd: SOp):
    super().__init__()

    if fst == snd:
      logging.warning("Creating a SequenceMap with both inputs being the same "
                      "SOp is discouraged. You should use a Map instead.")

    self.f = f
    self.fst = fst
    self.snd = snd
    assert isinstance(self.fst, SOp)
    assert isinstance(self.snd, SOp)
    assert callable(self.f) and not isinstance(self.f, RASPExpr)

  @property
  def children(self) -> Sequence[RASPExpr]:
    return [self.fst, self.snd]

class LinearSequenceMap(SequenceMap):
  """SOp that evaluates a linear function elementwise on the two given SOp's."""

  def __init__(self, fst: SOp, snd: SOp, fst_fac: float, snd_fac: float):
    super().__init__(fst=fst, snd=snd, f=lambda x, y: fst_fac * x + snd_fac * y)
    self.fst_fac = fst_fac
    self.snd_fac = snd_fac

class Full(SOp):
  """A SOp evaluating to [fill]*len(input_values)."""

  def __init__(self, fill: Value):
    super().__init__()
    self.fill = fill

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

def sop_not(sop: SOp) -> SOp:
  return Map(lambda t: not t, sop)

class ConstantSOp(SOp, Generic[VT]):
  """A constant S-Op for testing purposes."""

  def __init__(self, value: Sequence[VT], check_length: bool = True):
    super().__init__()
    self.value = value
    self.check_length = check_length

  @property
  def children(self) -> Sequence[RASPExpr]:
    return []

### Selectors.

class Predicate(Protocol):

  def __call__(self, key: Value, query: Value) -> bool:
    """Applies the predicate."""

class Comparison(enum.Enum):
  """A two-place boolean comparison predicate for use in Select."""
  EQ = "=="
  LT = "<"
  LEQ = "<="
  GT = ">"
  GEQ = ">="
  NEQ = "!="
  TRUE = "True"
  FALSE = "False"

  def __call__(self, key: Value, query: Value) -> bool:
    if key is None:
      raise ValueError("key is None!")
    if query is None:
      raise ValueError("query is None!")
    return _comparison_table[self](key, query)

_comparison_table = {
    Comparison.EQ: lambda key, query: key == query,
    Comparison.LT: lambda key, query: key < query,
    Comparison.LEQ: lambda key, query: key <= query,
    Comparison.GT: lambda key, query: key > query,
    Comparison.GEQ: lambda key, query: key >= query,
    Comparison.NEQ: lambda key, query: key != query,
    Comparison.TRUE: lambda key, query: True,
    Comparison.FALSE: lambda key, query: False,
}

class Selector(RASPExpr):
  """RASP Selector. Represents something like an attention head's weights."""

  def __call__(self, xs: Sequence[Value]) -> SelectorValue:
    return evaluate(self, xs)  # pytype: disable=bad-return-type

  # Allow construction of Selector combinations using Python logical operators.
  def __and__(self, other: "Selector") -> "Selector":
    """self & other."""
    return selector_and(self, other)

  def __rand__(self, other: "Selector") -> "Selector":
    """other & self."""
    return selector_and(other, self)

  def __or__(self, other: "Selector") -> "Selector":
    """self | other."""
    return selector_or(self, other)

  def __ror__(self, other: "Selector") -> "Selector":
    """other | self."""
    return selector_or(other, self)

  def __invert__(self) -> "Selector":
    """~self."""
    return selector_not(self)

class Select(Selector):
  """Primitive that creates a Selector."""

  def __init__(self, keys: SOp, queries: SOp, predicate: Predicate):

    super().__init__()
    self.keys = keys
    self.queries = queries
    self.predicate = predicate
    assert isinstance(self.keys, SOp)
    assert isinstance(self.queries, SOp)