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

from typing import List, Sequence

from tracr.rasp import rasp

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

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

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

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

length = make_length()

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

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

  Args:
    sop: an SOp

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

  Only supports unique keys.

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

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

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

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

  Note that this approach only works for numerical keys.

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

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

  Returns:
    Output SOp of sort program.
  """
  keys = rasp.SequenceMap(lambda x, i: x + min_key * i / max_seq_len, keys,
                          rasp.indices)
  return make_sort_unique(vals, keys)

def make_sort_freq(max_seq_len: int) -> rasp.SOp:
  """Returns tokens sorted by the frequency they appear in the input.

  Tokens the appear the same amount of times are output in the same order as in
  the input.

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

  Args:
    max_seq_len: Maximum sequence length (used to ensure keys are unique)
  """

  hist = -1 * make_hist().named("hist")
  return make_sort(
      rasp.tokens, hist, max_seq_len=max_seq_len, min_key=1).named("sort_freq")


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

from typing import List, Sequence

from tracr.rasp import rasp

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

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

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

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

length = make_length()

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

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

  Args:
    sop: an SOp

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

  Only supports unique keys.

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

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

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

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

  Note that this approach only works for numerical keys.

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

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

  Returns:
    Output SOp of sort program.
  """
  keys = rasp.SequenceMap(lambda x, i: x + min_key * i / max_seq_len, keys,
                          rasp.indices)
  return make_sort_unique(vals, keys)

def make_sort_freq(max_seq_len: int) -> rasp.SOp:
  """Returns tokens sorted by the frequency they appear in the input.

  Tokens the appear the same amount of times are output in the same order as in
  the input.

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

  Args:
    max_seq_len: Maximum sequence length (used to ensure keys are unique)
  """
  hist = -1 * make_hist().named("hist")
  return make_sort(
      rasp.tokens, hist, max_seq_len=max_seq_len, min_key=1).named("sort_freq")

### Programs that work under both causal and regular evaluation.

def make_frac_prevs(bools: rasp.SOp) -> rasp.SOp:
  """Count the fraction of previous tokens where a specific condition was True.

   (As implemented in the RASP paper.)

  Example usage:
    num_l = make_frac_prevs(rasp.tokens=="l")
    num_l("hello")
    >> [0, 0, 1/3, 1/2, 2/5]

  Args:
    bools: SOp mapping a sequence to a sequence of booleans.

  Returns:
    frac_prevs: SOp mapping an input to a sequence, where every element
      is the fraction of previous "True" tokens.
  """

  bools = rasp.numerical(bools)
  prevs = rasp.Select(rasp.indices, rasp.indices, rasp.Comparison.LEQ)
  return rasp.numerical(rasp.Aggregate(prevs, bools,
                                       default=0)).named("frac_prevs")


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

from typing import List, Sequence

from tracr.rasp import rasp

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

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

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

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

length = make_length()

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

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

  Args:
    sop: an SOp

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

  Only supports unique keys.

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

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

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

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

  Note that this approach only works for numerical keys.

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

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

  Returns:
    Output SOp of sort program.
  """
  keys = rasp.SequenceMap(lambda x, i: x + min_key * i / max_seq_len, keys,
                          rasp.indices)
  return make_sort_unique(vals, keys)

def make_sort_freq(max_seq_len: int) -> rasp.SOp:
  """Returns tokens sorted by the frequency they appear in the input.

  Tokens the appear the same amount of times are output in the same order as in
  the input.

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

  Args:
    max_seq_len: Maximum sequence length (used to ensure keys are unique)
  """
  hist = -1 * make_hist().named("hist")
  return make_sort(
      rasp.tokens, hist, max_seq_len=max_seq_len, min_key=1).named("sort_freq")

### Programs that work under both causal and regular evaluation.

def make_frac_prevs(bools: rasp.SOp) -> rasp.SOp:
  """Count the fraction of previous tokens where a specific condition was True.

   (As implemented in the RASP paper.)

  Example usage:
    num_l = make_frac_prevs(rasp.tokens=="l")
    num_l("hello")
    >> [0, 0, 1/3, 1/2, 2/5]

  Args:
    bools: SOp mapping a sequence to a sequence of booleans.

  Returns:
    frac_prevs: SOp mapping an input to a sequence, where every element
      is the fraction of previous "True" tokens.
  """
  bools = rasp.numerical(bools)
  prevs = rasp.Select(rasp.indices, rasp.indices, rasp.Comparison.LEQ)
  return rasp.numerical(rasp.Aggregate(prevs, bools,
                                       default=0)).named("frac_prevs")

def shift_by(offset: int, /, sop: rasp.SOp) -> rasp.SOp:
  """Returns the sop, shifted by `offset`, None-padded."""

  select_off_by_offset = rasp.Select(rasp.indices, rasp.indices,
                                     lambda k, q: q == k + offset)
  out = rasp.Aggregate(select_off_by_offset, sop, default=None)
  return out.named(f"shift_by({offset})")


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

from typing import List, Sequence

from tracr.rasp import rasp

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

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

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

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

length = make_length()

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

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

  Args:
    sop: an SOp

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

  Only supports unique keys.

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

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

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

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

  Note that this approach only works for numerical keys.

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

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

  Returns:
    Output SOp of sort program.
  """
  keys = rasp.SequenceMap(lambda x, i: x + min_key * i / max_seq_len, keys,
                          rasp.indices)
  return make_sort_unique(vals, keys)

def make_sort_freq(max_seq_len: int) -> rasp.SOp:
  """Returns tokens sorted by the frequency they appear in the input.

  Tokens the appear the same amount of times are output in the same order as in
  the input.

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

  Args:
    max_seq_len: Maximum sequence length (used to ensure keys are unique)
  """
  hist = -1 * make_hist().named("hist")
  return make_sort(
      rasp.tokens, hist, max_seq_len=max_seq_len, min_key=1).named("sort_freq")

### Programs that work under both causal and regular evaluation.

def make_frac_prevs(bools: rasp.SOp) -> rasp.SOp:
  """Count the fraction of previous tokens where a specific condition was True.

   (As implemented in the RASP paper.)

  Example usage:
    num_l = make_frac_prevs(rasp.tokens=="l")
    num_l("hello")
    >> [0, 0, 1/3, 1/2, 2/5]

  Args:
    bools: SOp mapping a sequence to a sequence of booleans.

  Returns:
    frac_prevs: SOp mapping an input to a sequence, where every element
      is the fraction of previous "True" tokens.
  """
  bools = rasp.numerical(bools)
  prevs = rasp.Select(rasp.indices, rasp.indices, rasp.Comparison.LEQ)
  return rasp.numerical(rasp.Aggregate(prevs, bools,
                                       default=0)).named("frac_prevs")

def shift_by(offset: int, /, sop: rasp.SOp) -> rasp.SOp:
  """Returns the sop, shifted by `offset`, None-padded."""
  select_off_by_offset = rasp.Select(rasp.indices, rasp.indices,
                                     lambda k, q: q == k + offset)
  out = rasp.Aggregate(select_off_by_offset, sop, default=None)
  return out.named(f"shift_by({offset})")

def detect_pattern(sop: rasp.SOp, pattern: Sequence[rasp.Value]) -> rasp.SOp:
  """Returns an SOp which is True at the final element of the pattern.

  The first len(pattern) - 1 elements of the output SOp are None-padded.

  detect_pattern(tokens, "abc")("abcabc") == [None, None, T, F, F, T]

  Args:
    sop: the SOp in which to look for patterns.
    pattern: a sequence of values to look for.

  Returns:
    a sop which detects the pattern.
  """

  if len(pattern) < 1:
    raise ValueError(f"Length of `pattern` must be at least 1. Got {pattern}")

  # detectors[i] will be a boolean-valued SOp which is true at position j iff
  # the i'th (from the end) element of the pattern was detected at position j-i.
  detectors = []
  for i, element in enumerate(reversed(pattern)):
    detector = sop == element
    if i != 0:
      detector = shift_by(i, detector)
    detectors.append(detector)

  # All that's left is to take the AND over all detectors.
  pattern_detected = detectors.pop()
  while detectors:
    pattern_detected = pattern_detected & detectors.pop()

  return pattern_detected.named(f"detect_pattern({pattern})")


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

from typing import List, Sequence

from tracr.rasp import rasp

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

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

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

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

length = make_length()

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

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

  Args:
    sop: an SOp

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

  Only supports unique keys.

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

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

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

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

  Note that this approach only works for numerical keys.

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

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

  Returns:
    Output SOp of sort program.
  """
  keys = rasp.SequenceMap(lambda x, i: x + min_key * i / max_seq_len, keys,
                          rasp.indices)
  return make_sort_unique(vals, keys)

def make_sort_freq(max_seq_len: int) -> rasp.SOp:
  """Returns tokens sorted by the frequency they appear in the input.

  Tokens the appear the same amount of times are output in the same order as in
  the input.

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

  Args:
    max_seq_len: Maximum sequence length (used to ensure keys are unique)
  """
  hist = -1 * make_hist().named("hist")
  return make_sort(
      rasp.tokens, hist, max_seq_len=max_seq_len, min_key=1).named("sort_freq")

### Programs that work under both causal and regular evaluation.

def make_frac_prevs(bools: rasp.SOp) -> rasp.SOp:
  """Count the fraction of previous tokens where a specific condition was True.

   (As implemented in the RASP paper.)

  Example usage:
    num_l = make_frac_prevs(rasp.tokens=="l")
    num_l("hello")
    >> [0, 0, 1/3, 1/2, 2/5]

  Args:
    bools: SOp mapping a sequence to a sequence of booleans.

  Returns:
    frac_prevs: SOp mapping an input to a sequence, where every element
      is the fraction of previous "True" tokens.
  """
  bools = rasp.numerical(bools)
  prevs = rasp.Select(rasp.indices, rasp.indices, rasp.Comparison.LEQ)
  return rasp.numerical(rasp.Aggregate(prevs, bools,
                                       default=0)).named("frac_prevs")

def shift_by(offset: int, /, sop: rasp.SOp) -> rasp.SOp:
  """Returns the sop, shifted by `offset`, None-padded."""
  select_off_by_offset = rasp.Select(rasp.indices, rasp.indices,
                                     lambda k, q: q == k + offset)
  out = rasp.Aggregate(select_off_by_offset, sop, default=None)
  return out.named(f"shift_by({offset})")

def detect_pattern(sop: rasp.SOp, pattern: Sequence[rasp.Value]) -> rasp.SOp:
  """Returns an SOp which is True at the final element of the pattern.

  The first len(pattern) - 1 elements of the output SOp are None-padded.

  detect_pattern(tokens, "abc")("abcabc") == [None, None, T, F, F, T]

  Args:
    sop: the SOp in which to look for patterns.
    pattern: a sequence of values to look for.

  Returns:
    a sop which detects the pattern.
  """

  if len(pattern) < 1:
    raise ValueError(f"Length of `pattern` must be at least 1. Got {pattern}")

  # detectors[i] will be a boolean-valued SOp which is true at position j iff
  # the i'th (from the end) element of the pattern was detected at position j-i.
  detectors = []
  for i, element in enumerate(reversed(pattern)):
    detector = sop == element
    if i != 0:
      detector = shift_by(i, detector)
    detectors.append(detector)

  # All that's left is to take the AND over all detectors.
  pattern_detected = detectors.pop()
  while detectors:
    pattern_detected = pattern_detected & detectors.pop()

  return pattern_detected.named(f"detect_pattern({pattern})")

def make_count_less_freq(n: int) -> rasp.SOp:
  """Returns how many tokens appear fewer than n times in the input.

  The output sequence contains this count in each position.

  Example usage:
    count_less_freq = make_count_less_freq(2)
    count_less_freq(["a", "a", "a", "b", "b", "c"])
    >> [3, 3, 3, 3, 3, 3]
    count_less_freq(["a", "a", "c", "b", "b", "c"])
    >> [6, 6, 6, 6, 6, 6]

  Args:
    n: Integer to compare token frequences to.
  """

  hist = make_hist().named("hist")
  select_less = rasp.Select(hist, hist,
                            lambda x, y: x <= n).named("select_less")
  return rasp.SelectorWidth(select_less).named("count_less_freq")


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

from typing import List, Sequence

from tracr.rasp import rasp

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

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

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

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

length = make_length()

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

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

  Args:
    sop: an SOp

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

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

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

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

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

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

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

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

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

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

   (As implemented in the RASP paper.)

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

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

  Only supports unique keys.

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

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

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

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

  Note that this approach only works for numerical keys.

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

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

  Returns:
    Output SOp of sort program.
  """
  keys = rasp.SequenceMap(lambda x, i: x + min_key * i / max_seq_len, keys,
                          rasp.indices)
  return make_sort_unique(vals, keys)

def make_sort_freq(max_seq_len: int) -> rasp.SOp:
  """Returns tokens sorted by the frequency they appear in the input.

  Tokens the appear the same amount of times are output in the same order as in
  the input.

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

  Args:
    max_seq_len: Maximum sequence length (used to ensure keys are unique)
  """
  hist = -1 * make_hist().named("hist")
  return make_sort(
      rasp.tokens, hist, max_seq_len=max_seq_len, min_key=1).named("sort_freq")

### Programs that work under both causal and regular evaluation.

def make_frac_prevs(bools: rasp.SOp) -> rasp.SOp:
  """Count the fraction of previous tokens where a specific condition was True.

   (As implemented in the RASP paper.)

  Example usage:
    num_l = make_frac_prevs(rasp.tokens=="l")
    num_l("hello")
    >> [0, 0, 1/3, 1/2, 2/5]

  Args:
    bools: SOp mapping a sequence to a sequence of booleans.

  Returns:
    frac_prevs: SOp mapping an input to a sequence, where every element
      is the fraction of previous "True" tokens.
  """
  bools = rasp.numerical(bools)
  prevs = rasp.Select(rasp.indices, rasp.indices, rasp.Comparison.LEQ)
  return rasp.numerical(rasp.Aggregate(prevs, bools,
                                       default=0)).named("frac_prevs")

def shift_by(offset: int, /, sop: rasp.SOp) -> rasp.SOp:
  """Returns the sop, shifted by `offset`, None-padded."""
  select_off_by_offset = rasp.Select(rasp.indices, rasp.indices,
                                     lambda k, q: q == k + offset)
  out = rasp.Aggregate(select_off_by_offset, sop, default=None)
  return out.named(f"shift_by({offset})")

def detect_pattern(sop: rasp.SOp, pattern: Sequence[rasp.Value]) -> rasp.SOp:
  """Returns an SOp which is True at the final element of the pattern.

  The first len(pattern) - 1 elements of the output SOp are None-padded.

  detect_pattern(tokens, "abc")("abcabc") == [None, None, T, F, F, T]

  Args:
    sop: the SOp in which to look for patterns.
    pattern: a sequence of values to look for.

  Returns:
    a sop which detects the pattern.
  """

  if len(pattern) < 1:
    raise ValueError(f"Length of `pattern` must be at least 1. Got {pattern}")

  # detectors[i] will be a boolean-valued SOp which is true at position j iff
  # the i'th (from the end) element of the pattern was detected at position j-i.
  detectors = []
  for i, element in enumerate(reversed(pattern)):
    detector = sop == element
    if i != 0:
      detector = shift_by(i, detector)
    detectors.append(detector)

  # All that's left is to take the AND over all detectors.
  pattern_detected = detectors.pop()
  while detectors:
    pattern_detected = pattern_detected & detectors.pop()

  return pattern_detected.named(f"detect_pattern({pattern})")

def make_count_less_freq(n: int) -> rasp.SOp:
  """Returns how many tokens appear fewer than n times in the input.

  The output sequence contains this count in each position.

  Example usage:
    count_less_freq = make_count_less_freq(2)
    count_less_freq(["a", "a", "a", "b", "b", "c"])
    >> [3, 3, 3, 3, 3, 3]
    count_less_freq(["a", "a", "c", "b", "b", "c"])
    >> [6, 6, 6, 6, 6, 6]

  Args:
    n: Integer to compare token frequences to.
  """
  hist = make_hist().named("hist")
  select_less = rasp.Select(hist, hist,
                            lambda x, y: x <= n).named("select_less")
  return rasp.SelectorWidth(select_less).named("count_less_freq")

def make_count(sop, token):
  """Returns the count of `token` in `sop`.

  The output sequence contains this count in each position.

  Example usage:
    count = make_count(tokens, "a")
    count(["a", "a", "a", "b", "b", "c"])
    >> [3, 3, 3, 3, 3, 3]
    count(["c", "a", "b", "c"])
    >> [1, 1, 1, 1]

  Args:
    sop: Sop to count tokens in.
    token: Token to count.
  """
  return rasp.SelectorWidth(rasp.Select(
      sop, sop, lambda k, q: k == token)).named(f"count_{token}")

def make_nary_sequencemap(f, *sops):
  """Returns an SOp that simulates an n-ary SequenceMap.

  Uses multiple binary SequenceMaps to convert n SOps x_1, x_2, ..., x_n
  into a single SOp arguments that takes n-tuples as value. The n-ary sequence
  map implementing f is then a Map on this resulting SOp.

  Note that the intermediate variables representing tuples of varying length
  will be encoded categorically, and can become very high-dimensional. So,
  using this function might lead to very large compiled models.

  Args:
    f: Function with n arguments.
    *sops: Sequence of SOps, one for each argument of f.
  """

  values, *sops = sops
  for sop in sops:
    # x is a single entry in the first iteration but a tuple in later iterations
    values = rasp.SequenceMap(
        lambda x, y: (*x, y) if isinstance(x, tuple) else (x, y), values, sop)
  return rasp.Map(lambda args: f(*args), values)


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

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.compiler import assemble
from tracr.craft import bases

class AssembleTest(parameterized.TestCase):

  def test_token_embedding_produces_correct_embedding(self):
    # Token embeddings should be one-hot embeddings of the input integers
    # into the token subspace of residual_space
    input_space = bases.VectorSpaceWithBasis.from_values("0inp", range(2))
    indices_space = bases.VectorSpaceWithBasis.from_values("1ind", range(3))
    output_space = bases.VectorSpaceWithBasis.from_values("2out", range(2))
    residual_space = bases.join_vector_spaces(input_space, indices_space,
                                              output_space)

    @hk.without_apply_rng
    @hk.transform
    def token_pos_embed(tokens):

      embed_modules = assemble._make_embedding_modules(
          residual_space=residual_space,
          tokens_space=input_space,
          indices_space=indices_space,
          output_space=output_space)
      return embed_modules.token_embed(tokens)


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

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.compiler import assemble
from tracr.craft import bases

class AssembleTest(parameterized.TestCase):

  def test_token_embedding_produces_correct_embedding(self):
    # Token embeddings should be one-hot embeddings of the input integers
    # into the token subspace of residual_space
    input_space = bases.VectorSpaceWithBasis.from_values("0inp", range(2))
    indices_space = bases.VectorSpaceWithBasis.from_values("1ind", range(3))
    output_space = bases.VectorSpaceWithBasis.from_values("2out", range(2))
    residual_space = bases.join_vector_spaces(input_space, indices_space,
                                              output_space)

    @hk.without_apply_rng
    @hk.transform
    def token_pos_embed(tokens):
      embed_modules = assemble._make_embedding_modules(
          residual_space=residual_space,
          tokens_space=input_space,
          indices_space=indices_space,
          output_space=output_space)
      return embed_modules.token_embed(tokens)

    tokens = jnp.array([0, 0, 1])
    expected_token_embeddings = jnp.array([[1, 0, 0, 0, 0, 0, 0],
                                           [1, 0, 0, 0, 0, 0, 0],
                                           [0, 1, 0, 0, 0, 0, 0]])

    params = token_pos_embed.init(jax.random.PRNGKey(0), tokens)
    embeddings = token_pos_embed.apply(params, tokens)
    np.testing.assert_allclose(embeddings, expected_token_embeddings)

  def test_position_embedding_produces_correct_embedding(self):
    # Position embeddings should be one-hot embeddings of the input integers
    # (representing indices) into the indices subspace of residual_space
    input_space = bases.VectorSpaceWithBasis.from_values("0inp", range(2))
    indices_space = bases.VectorSpaceWithBasis.from_values("1ind", range(3))
    output_space = bases.VectorSpaceWithBasis.from_values("2out", range(2))
    residual_space = bases.join_vector_spaces(input_space, indices_space,
                                              output_space)

    @hk.without_apply_rng
    @hk.transform
    def token_pos_embed(tokens):

      embed_modules = assemble._make_embedding_modules(
          residual_space=residual_space,
          tokens_space=input_space,
          indices_space=indices_space,
          output_space=output_space)
      return embed_modules.pos_embed(jnp.indices(tokens.shape)[-1])


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

from absl.testing import absltest
from absl.testing import parameterized
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from tracr.compiler import assemble
from tracr.craft import bases

class AssembleTest(parameterized.TestCase):

  def test_token_embedding_produces_correct_embedding(self):
    # Token embeddings should be one-hot embeddings of the input integers
    # into the token subspace of residual_space
    input_space = bases.VectorSpaceWithBasis.from_values("0inp", range(2))
    indices_space = bases.VectorSpaceWithBasis.from_values("1ind", range(3))
    output_space = bases.VectorSpaceWithBasis.from_values("2out", range(2))
    residual_space = bases.join_vector_spaces(input_space, indices_space,
                                              output_space)

    @hk.without_apply_rng
    @hk.transform
    def token_pos_embed(tokens):
      embed_modules = assemble._make_embedding_modules(
          residual_space=residual_space,
          tokens_space=input_space,
          indices_space=indices_space,
          output_space=output_space)
      return embed_modules.token_embed(tokens)

    tokens = jnp.array([0, 0, 1])
    expected_token_embeddings = jnp.array([[1, 0, 0, 0, 0, 0, 0],
                                           [1, 0, 0, 0, 0, 0, 0],
                                           [0, 1, 0, 0, 0, 0, 0]])

    params = token_pos_embed.init(jax.random.PRNGKey(0), tokens)
    embeddings = token_pos_embed.apply(params, tokens)
    np.testing.assert_allclose(embeddings, expected_token_embeddings)

  def test_position_embedding_produces_correct_embedding(self):
    # Position embeddings should be one-hot embeddings of the input integers
    # (representing indices) into the indices subspace of residual_space
    input_space = bases.VectorSpaceWithBasis.from_values("0inp", range(2))
    indices_space = bases.VectorSpaceWithBasis.from_values("1ind", range(3))
    output_space = bases.VectorSpaceWithBasis.from_values("2out", range(2))
    residual_space = bases.join_vector_spaces(input_space, indices_space,
                                              output_space)

    @hk.without_apply_rng
    @hk.transform
    def token_pos_embed(tokens):
      embed_modules = assemble._make_embedding_modules(
          residual_space=residual_space,
          tokens_space=input_space,
          indices_space=indices_space,
          output_space=output_space)
      return embed_modules.pos_embed(jnp.indices(tokens.shape)[-1])

    tokens = jnp.array([3, 0, 0, 1])
    expected_pos_embeddings = jnp.array([[0, 0, 0, 0, 0, 0, 0],
                                         [0, 0, 1, 0, 0, 0, 0],
                                         [0, 0, 0, 1, 0, 0, 0],
                                         [0, 0, 0, 0, 1, 0, 0]])

    params = token_pos_embed.init(jax.random.PRNGKey(0), tokens)
    embeddings = token_pos_embed.apply(params, tokens)
    np.testing.assert_allclose(embeddings, expected_pos_embeddings)

  def test_unembedding(self):
    # Prepend numbers to preserve basis order [input, index, output]
    input_space = bases.VectorSpaceWithBasis.from_values("0inp", range(2))
    indices_space = bases.VectorSpaceWithBasis.from_values("1ind", range(3))
    output_space = bases.VectorSpaceWithBasis.from_values("2out", range(2))
    residual_space = bases.join_vector_spaces(input_space, indices_space,
                                              output_space)

    @hk.without_apply_rng
    @hk.transform
    def unembed(embeddings):

      embed_modules = assemble._make_embedding_modules(
          residual_space=residual_space,
          tokens_space=input_space,
          indices_space=indices_space,
          output_space=output_space)
      return embed_modules.unembed(embeddings, use_unembed_argmax=True)


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

from typing import Optional

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

class QueryKeyToAttnLogit(Protocol):

  def __call__(self, query: bases.BasisDirection,
               key: bases.BasisDirection) -> bool:
    pass

def categorical_attn(
    query_space: bases.VectorSpaceWithBasis,
    key_space: bases.VectorSpaceWithBasis,
    value_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    bos_space: bases.VectorSpaceWithBasis,
    one_space: bases.VectorSpaceWithBasis,
    attn_fn: QueryKeyToAttnLogit,
    default_output: Optional[bases.VectorInBasis] = None,
    causal: bool = False,
    always_attend_to_bos: bool = False,
    use_bos_for_default_output: bool = True,
    softmax_coldness: float = 100.,
) -> transformers.AttentionHead:
  """Returns an attention head for categorical inputs.

  Assumes the existence of a beginning of sequence token and attends to it
  always with strength 0.5*softmax_coldness. This allows to implement an
  arbitrary default value for rows in the attention pattern that are all-zero.

  Attends to the BOS token if all other key-query pairs have zero attention.
  Hence, the first value in the value sequence will be the default output for
  such cases.

  Args:
    query_space: Vector space containing (categorical) query input.
    key_space: Vector space containing (categorical) key input.
    value_space: Vector space containing (numerical) value input.
    output_space: Vector space which will contain (numerical) output.
    bos_space: 1-d space used to identify the beginning of sequence token.
    one_space: 1-d space which contains 1 at every position.
    attn_fn: A selector function f(query, key) operating on the query/key basis
      directions that defines the attention pattern.
    default_output: Output to return if attention pattern is all zero.
    causal: If True, use masked attention.
    always_attend_to_bos: If True, always attend to the BOS token. If False,
      only attend to BOS when attending to nothing else.
    use_bos_for_default_output: If True, assume BOS is not in the value space
      and output a default value when attending to BOS. If False, assume BOS is
      in the value space, and map it to the output space like any other token.
    softmax_coldness: The inverse temperature of the softmax. Default value is
      high which makes the attention close to a hard maximum.
  """
  bases.ensure_dims(bos_space, num_dims=1, name="bos_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")
  bos_direction = bos_space.basis[0]
  one_direction = one_space.basis[0]

  # Add bos direction to query, key, and value spaces in case it is missing
  query_space = bases.join_vector_spaces(query_space, bos_space, one_space)
  key_space = bases.join_vector_spaces(key_space, bos_space)
  value_space = bases.join_vector_spaces(value_space, bos_space)

  if always_attend_to_bos:
    value_basis = value_space.basis
  else:
    value_basis = [v for v in value_space.basis if v != bos_direction]
  assert len(value_basis) == output_space.num_dims
  value_to_output = dict(zip(value_basis, output_space.basis))

  if default_output is None:
    default_output = output_space.null_vector()
  assert default_output in output_space

  def qk_fun(query: bases.BasisDirection, key: bases.BasisDirection) -> float:

    # We want to enforce the following property on our attention patterns:
    # - if nothing else is attended to, attend to the BOS token.
    # - otherwise, don't attend to the BOS token.
    #
    # We assume that the BOS position always only contains the vector bos + one,
    # and that any other position has bos coefficient 0.
    #
    # We do this as follows:
    # Let Q and K be subspaces of V containing the query and key vectors,
    # both disjoint with the BOS space {bos} or the one space {one}.
    # Suppose we have an attn_fn which defines a bilinear W_QK: V x V -> ℝ,
    # s.t. W_QK(q, k) = 0 whenever either q or k are bos or one.
    #
    # Then define W_new: V x V -> ℝ st:
    # W_new(one, bos) = 0.5, otherwise 0.
    #
    # Now set W_QK' = W_QK + W_new.
    #
    # To evaluate the attention to the BOS position:
    # W_QK'(q, bos + one)
    # = W_QK'(q, bos) + W_QK'(q, one)
    # = W_QK(q, bos) + W_QK(q, one) + W_new(q, bos) + W_new(q, one)
    # = 0            + 0            + W_new(q, bos) + W_new(q, one)
    # = W_new(q, bos) + W_new(q, one)
    # = W_new(q' + one, bos) + W_new(q' + one, one)  where q = one + q'
    # = W_new(q', bos) + W_new(one, bos) + W_new(q', one) + W_new(one, one)
    # = 0              + 0.5             + 0              + 0
    # = 0.5
    #
    # To evaluate the attention to a non-BOS position:
    # W_QK'(0 * bos + q, 0 * bos + k)  # s.t. q ∈ Q+{one}, k ∈ K+{one}
    # = 0*W_QK'(bos, 0*bos + k) + W_QK'(q, 0*bos + k)
    # = W_QK'(q, 0*bos + k)
    # = 0*W_QK'(q, bos) + W_QK'(q, k)
    # = W_QK'(q, k)
    # = W_QK(q, k)    since W_QK' = W_QK on inputs not containing bos.
    # = W_QK(q', k')  since W_QK(x, y) = 0 whenever x or y are one.
    #
    # Since W_QK(q, k) takes values in 0, 1, a sufficiently high softmax
    # coldness will give us the desired property.                            QED
    #
    # The following implements this idea.
    # By replacing 0.5 with 1, we can instead enforce a different property: that
    # the BOS token is always attended to in addition to whatever else.

    if key == bos_direction and query == one_direction:
      c = 1. if always_attend_to_bos else 0.5
      return c * softmax_coldness
    elif {key, query}.intersection({one_direction, bos_direction}):
      return 0

    return softmax_coldness * attn_fn(query, key)


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.
# ==============================================================================
"""Attention head for categorical inputs."""

from typing import Optional

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

class QueryKeyToAttnLogit(Protocol):

  def __call__(self, query: bases.BasisDirection,
               key: bases.BasisDirection) -> bool:
    pass

def categorical_attn(
    query_space: bases.VectorSpaceWithBasis,
    key_space: bases.VectorSpaceWithBasis,
    value_space: bases.VectorSpaceWithBasis,
    output_space: bases.VectorSpaceWithBasis,
    bos_space: bases.VectorSpaceWithBasis,
    one_space: bases.VectorSpaceWithBasis,
    attn_fn: QueryKeyToAttnLogit,
    default_output: Optional[bases.VectorInBasis] = None,
    causal: bool = False,
    always_attend_to_bos: bool = False,
    use_bos_for_default_output: bool = True,
    softmax_coldness: float = 100.,
) -> transformers.AttentionHead:
  """Returns an attention head for categorical inputs.

  Assumes the existence of a beginning of sequence token and attends to it
  always with strength 0.5*softmax_coldness. This allows to implement an
  arbitrary default value for rows in the attention pattern that are all-zero.

  Attends to the BOS token if all other key-query pairs have zero attention.
  Hence, the first value in the value sequence will be the default output for
  such cases.

  Args:
    query_space: Vector space containing (categorical) query input.
    key_space: Vector space containing (categorical) key input.
    value_space: Vector space containing (numerical) value input.
    output_space: Vector space which will contain (numerical) output.
    bos_space: 1-d space used to identify the beginning of sequence token.
    one_space: 1-d space which contains 1 at every position.
    attn_fn: A selector function f(query, key) operating on the query/key basis
      directions that defines the attention pattern.
    default_output: Output to return if attention pattern is all zero.
    causal: If True, use masked attention.
    always_attend_to_bos: If True, always attend to the BOS token. If False,
      only attend to BOS when attending to nothing else.
    use_bos_for_default_output: If True, assume BOS is not in the value space
      and output a default value when attending to BOS. If False, assume BOS is
      in the value space, and map it to the output space like any other token.
    softmax_coldness: The inverse temperature of the softmax. Default value is
      high which makes the attention close to a hard maximum.
  """
  bases.ensure_dims(bos_space, num_dims=1, name="bos_space")
  bases.ensure_dims(one_space, num_dims=1, name="one_space")
  bos_direction = bos_space.basis[0]
  one_direction = one_space.basis[0]

  # Add bos direction to query, key, and value spaces in case it is missing
  query_space = bases.join_vector_spaces(query_space, bos_space, one_space)
  key_space = bases.join_vector_spaces(key_space, bos_space)
  value_space = bases.join_vector_spaces(value_space, bos_space)

  if always_attend_to_bos:
    value_basis = value_space.basis
  else:
    value_basis = [v for v in value_space.basis if v != bos_direction]
  assert len(value_basis) == output_space.num_dims
  value_to_output = dict(zip(value_basis, output_space.basis))

  if default_output is None:
    default_output = output_space.null_vector()
  assert default_output in output_space

  def qk_fun(query: bases.BasisDirection, key: bases.BasisDirection) -> float:

    # We want to enforce the following property on our attention patterns:
    # - if nothing else is attended to, attend to the BOS token.
    # - otherwise, don't attend to the BOS token.
    #
    # We assume that the BOS position always only contains the vector bos + one,
    # and that any other position has bos coefficient 0.
    #
    # We do this as follows:
    # Let Q and K be subspaces of V containing the query and key vectors,
    # both disjoint with the BOS space {bos} or the one space {one}.
    # Suppose we have an attn_fn which defines a bilinear W_QK: V x V -> ℝ,
    # s.t. W_QK(q, k) = 0 whenever either q or k are bos or one.
    #
    # Then define W_new: V x V -> ℝ st:
    # W_new(one, bos) = 0.5, otherwise 0.
    #
    # Now set W_QK' = W_QK + W_new.
    #
    # To evaluate the attention to the BOS position:
    # W_QK'(q, bos + one)
    # = W_QK'(q, bos) + W_QK'(q, one)
    # = W_QK(q, bos) + W_QK(q, one) + W_new(q, bos) + W_new(q, one)
    # = 0            + 0            + W_new(q, bos) + W_new(q, one)
    # = W_new(q, bos) + W_new(q, one)
    # = W_new(q' + one, bos) + W_new(q' + one, one)  where q = one + q'
    # = W_new(q', bos) + W_new(one, bos) + W_new(q', one) + W_new(one, one)
    # = 0              + 0.5             + 0              + 0
    # = 0.5
    #
    # To evaluate the attention to a non-BOS position:
    # W_QK'(0 * bos + q, 0 * bos + k)  # s.t. q ∈ Q+{one}, k ∈ K+{one}
    # = 0*W_QK'(bos, 0*bos + k) + W_QK'(q, 0*bos + k)
    # = W_QK'(q, 0*bos + k)
    # = 0*W_QK'(q, bos) + W_QK'(q, k)
    # = W_QK'(q, k)
    # = W_QK(q, k)    since W_QK' = W_QK on inputs not containing bos.
    # = W_QK(q', k')  since W_QK(x, y) = 0 whenever x or y are one.
    #
    # Since W_QK(q, k) takes values in 0, 1, a sufficiently high softmax
    # coldness will give us the desired property.                            QED
    #
    # The following implements this idea.
    # By replacing 0.5 with 1, we can instead enforce a different property: that
    # the BOS token is always attended to in addition to whatever else.

    if key == bos_direction and query == one_direction:
      c = 1. if always_attend_to_bos else 0.5
      return c * softmax_coldness
    elif {key, query}.intersection({one_direction, bos_direction}):
      return 0

    return softmax_coldness * attn_fn(query, key)

  w_qk = vectorspace_fns.ScalarBilinear.from_action(
      query_space,
      key_space,
      qk_fun,
  )

  def ov_fun(input_dir: bases.BasisDirection) -> bases.VectorInBasis:

    if use_bos_for_default_output and input_dir == bos_direction:
      return default_output
    return output_space.vector_from_basis_direction(value_to_output[input_dir])


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

import dataclasses

from typing import Callable, Iterable, List

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

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

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

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

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

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

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

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

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

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

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

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

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

  def action(direction: bases.BasisDirection) -> bases.VectorInBasis:
    # hidden_k_0 = ReLU(L * (x - threshold[k]) + 1)
    # hidden_k_1 = ReLU(L * (x - threshold[k]))

    if direction == one_direction:
      hidden = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(f"{hidden_name}start"))
    else:
      hidden = hidden_space.null_vector()
    for k in range(1, num_vals):
      vec0 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 0)))
      vec1 = hidden_space.vector_from_basis_direction(
          bases.BasisDirection(hidden_name, (k, 1)))
      if direction == one_direction:
        hidden += (1 - large_number * value_thresholds[k - 1]) * vec0
        hidden -= large_number * value_thresholds[k - 1] * vec1
      else:
        hidden += large_number * vec0 + large_number * vec1
    return hidden


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

import dataclasses

from typing import Callable, Iterable, List

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  def second_layer_action(
      direction: bases.BasisDirection) -> bases.VectorInBasis:
    # output = sum(
    #     (hidden_k_0 - hidden_k_1) * (f(input[k]) - f(input[k-1]))
    #   for all k)

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