-- | This module provides the core of the RASP-L language.
--
-- It is based on Listing 2 of
-- "What Algorithms Can Transformers Learn", https://arxiv.org/abs/2310.16028,
-- by Zhou et al.
module RaskellCore
  ( -- * Types

    -- | Most of these types are merely aliases.

    -- ** Tokens and Sequences
    Token,
    Sequence,
    Keys,
    Queries,
    Values,

    -- ** Predicates and Selectors
    Predicate,
    Selector,
    BoolSequence,

    -- ** Aggregation
    Aggregator,
    AggregationType (..),

    -- * Functions

    -- ** Key-Query-Value lookup
    kqv,
    maxKQV,
    minKQV,
    selectCausal,
    filledWith,
    indicesOf,

    -- ** Aggregation
    aggregate,
    aggrMax,
    aggrMean,
    aggrMin,

    -- ** Selection
    selWidth,
    fromBool,

    -- * Compatibility

    -- | These functions are provided for closer compatibility with the original Python implementation of RASP-L.
    tokMap,
    seqMap,
    full,
    indices,
    aggr,
    select,
  )
where

import Data.Int (Int8)
import Data.Maybe (fromMaybe)

-- | A `Token` in a `Sequence` is a small integer.
-- RASP-L uses `Int8` to ensure all maps of type `Token` -> `Token` are learnable.
type Token = Int8

-- | A `Sequence` is a list of `Token`s.
type Sequence = [Token]

-- | A collection of keys is a list of `Token`s.
type Keys = Sequence

-- | A collection of queries is a list of `Token`s.
type Queries = Sequence

-- | A collection of values is a list of `Token`s.
type Values = Sequence

-- | We can compare `Keys` and `Queries` to determine if they match.
type Predicate = Token -> Token -> Bool

-- | The equivalents of "attention maps" are collections of `Bool`ean sequences.
type Selector = [BoolSequence]

-- | Internally, we sometimes need to operate on collections of `Bool`s.
type BoolSequence = [Bool]

-- | Enum for the three methods for aggregating selected values
data AggregationType
  = Min
  | Mean
  | Max

-- | Performs a key-query-value lookup operation and aggregates over values.
--
-- Given a filler token, an aggregation type, two sequences (keys and queries),
-- and a predicate, it returns a processed sequence. It first selects elements
-- based on the predicate and then aggregates them.
--
-- Roughly matches the attention layer of a Transformer.
kqv ::
  -- | Filler token used in aggregation
  Token ->
  -- | Type of aggregation (Min, Mean, Max)
  AggregationType ->
  -- | Sequence of keys
  Keys ->
  -- | Sequence of queries
  Queries ->
  -- | A boolean predicate that determines whether a key and query match
  Predicate ->
  -- | Sequence of values
  Values ->
  -- | The output sequence
  Sequence
kqv :: Token
-> AggregationType -> Keys -> Keys -> Predicate -> Keys -> Keys
kqv Token
filler AggregationType
agg Keys
keys Keys
queries Predicate
predicate = AggregationType -> Token -> Aggregator
aggregate AggregationType
agg Token
filler Aggregator -> Aggregator
forall a b. (a -> b) -> a -> b
$ Keys -> Keys -> Predicate -> Selector
selectCausal Keys
keys Keys
queries Predicate
predicate

-- | Performs Key-Query-Value lookup with maximum aggregation of values.
maxKQV :: Keys -> Queries -> Predicate -> Values -> Sequence
maxKQV :: Keys -> Keys -> Predicate -> Keys -> Keys
maxKQV = Token
-> AggregationType -> Keys -> Keys -> Predicate -> Keys -> Keys
kqv Token
minInt8 AggregationType
Max
  where
    minInt8 :: Token
minInt8 = Token
forall a. Bounded a => a
minBound :: Int8

-- | Performs Key-Query-Value lookup with minimum aggregation of values.
minKQV :: Keys -> Queries -> Predicate -> Values -> Sequence
minKQV :: Keys -> Keys -> Predicate -> Keys -> Keys
minKQV = Token
-> AggregationType -> Keys -> Keys -> Predicate -> Keys -> Keys
kqv Token
maxInt8 AggregationType
Min
  where
    maxInt8 :: Token
maxInt8 = Token
forall a. Bounded a => a
maxBound :: Int8

-- | Compares pairs of elements from sequences with a predicate subject to a causal constraint.
selectCausal :: Keys -> Queries -> Predicate -> Selector
selectCausal :: Keys -> Keys -> Predicate -> Selector
selectCausal Keys
keys Keys
queries Predicate
predicate =
  [ [ (Int
keyIndex Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
queryIndex) Bool -> Bool -> Bool
&& Predicate
predicate (Keys
keys Keys -> Int -> Token
forall a. HasCallStack => [a] -> Int -> a
!! Int
keyIndex) (Keys
queries Keys -> Int -> Token
forall a. HasCallStack => [a] -> Int -> a
!! Int
queryIndex)
      | Int
keyIndex <- [Int
0 .. Keys -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Keys
keys Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
    ]
    | Int
queryIndex <- [Int
0 .. Keys -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Keys
queries Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
  ]

-- | Creates a matched-length constant sequence with the provided token.
filledWith :: Sequence -> Token -> Sequence
filledWith :: Keys -> Token -> Keys
filledWith = Int -> Token -> Keys
forall a. Int -> a -> [a]
replicate (Int -> Token -> Keys) -> (Keys -> Int) -> Keys -> Token -> Keys
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Keys -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length

-- | Extracts the indices of the elements in a sequence.
indicesOf :: Sequence -> Sequence
indicesOf :: Keys -> Keys
indicesOf Keys
x = [Token
0 .. (Int -> Token
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Keys -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Keys
x) Token -> Token -> Token
forall a. Num a => a -> a -> a
- Token
1)]

-- | Type alias for "fully-specified" aggregators that are ready to aggregate a sequence of values with a selector.
type Aggregator = Selector -> Values -> Sequence

-- | Aggregates values with some aggregation, filling in with a default token.
aggregate :: AggregationType -> Token -> Aggregator
aggregate :: AggregationType -> Token -> Aggregator
aggregate AggregationType
Max = Token -> Aggregator
aggrMax
aggregate AggregationType
Mean = Token -> Aggregator
aggrMean
aggregate AggregationType
Min = Token -> Aggregator
aggrMin

-- | Aggregates values by selecting the largest value.
aggrMax :: Token -> Aggregator
aggrMax :: Token -> Aggregator
aggrMax Token
filler Selector
a Keys
v = ([Bool] -> Token) -> Selector -> Keys
forall a b. (a -> b) -> [a] -> [b]
map (Token -> Keys -> [Bool] -> Token
aggrMaxByRow Token
filler Keys
v) Selector
a

-- | Aggregates values by taking the mean.
aggrMean :: Token -> Aggregator
aggrMean :: Token -> Aggregator
aggrMean Token
filler Selector
a Keys
v = ([Bool] -> Token) -> Selector -> Keys
forall a b. (a -> b) -> [a] -> [b]
map (Token -> Keys -> [Bool] -> Token
aggrMeanByRow Token
filler Keys
v) Selector
a

-- | Aggregates values by selecting the smallest value.
aggrMin :: Token -> Aggregator
aggrMin :: Token -> Aggregator
aggrMin Token
filler Selector
a Keys
v = ([Bool] -> Token) -> Selector -> Keys
forall a b. (a -> b) -> [a] -> [b]
map (Token -> Keys -> [Bool] -> Token
aggrMinByRow Token
filler Keys
v) Selector
a

aggrMaxByRow :: Token -> Sequence -> BoolSequence -> Token
aggrMaxByRow :: Token -> Keys -> [Bool] -> Token
aggrMaxByRow Token
filler Keys
v [Bool]
a = Token -> Maybe Token -> Token
forall a. a -> Maybe a -> a
fromMaybe Token
filler Maybe Token
maybeMax
  where
    maybeMax :: Maybe Token
maybeMax = Keys -> Maybe Token
forall a. Ord a => [a] -> Maybe a
safeMaximum ([Bool] -> Keys -> Keys
forall a. [Bool] -> [a] -> [a]
filterByList [Bool]
a Keys
v)

aggrMeanByRow :: Token -> Sequence -> BoolSequence -> Token
aggrMeanByRow :: Token -> Keys -> [Bool] -> Token
aggrMeanByRow Token
filler Keys
v [Bool]
a = Token -> Maybe Token -> Token
forall a. a -> Maybe a -> a
fromMaybe Token
filler Maybe Token
maybeMean
  where
    maybeMean :: Maybe Token
maybeMean = Keys -> Maybe Token
safeInt8Mean ([Bool] -> Keys -> Keys
forall a. [Bool] -> [a] -> [a]
filterByList [Bool]
a Keys
v)

aggrMinByRow :: Token -> Sequence -> BoolSequence -> Token
aggrMinByRow :: Token -> Keys -> [Bool] -> Token
aggrMinByRow Token
filler Keys
v [Bool]
a = Token -> Maybe Token -> Token
forall a. a -> Maybe a -> a
fromMaybe Token
filler Maybe Token
maybeMin
  where
    maybeMin :: Maybe Token
maybeMin = Keys -> Maybe Token
forall a. Ord a => [a] -> Maybe a
safeMinimum ([Bool] -> Keys -> Keys
forall a. [Bool] -> [a] -> [a]
filterByList [Bool]
a Keys
v)

filterByList :: [Bool] -> [a] -> [a]
filterByList :: forall a. [Bool] -> [a] -> [a]
filterByList (Bool
True : [Bool]
bs) (a
x : [a]
xs) = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [Bool] -> [a] -> [a]
forall a. [Bool] -> [a] -> [a]
filterByList [Bool]
bs [a]
xs
filterByList (Bool
False : [Bool]
bs) (a
_ : [a]
xs) = [Bool] -> [a] -> [a]
forall a. [Bool] -> [a] -> [a]
filterByList [Bool]
bs [a]
xs
filterByList [Bool]
_ [a]
_ = []

safeMaximum :: (Ord a) => [a] -> Maybe a
safeMaximum :: forall a. Ord a => [a] -> Maybe a
safeMaximum [] = Maybe a
forall a. Maybe a
Nothing
safeMaximum [a]
xs = a -> Maybe a
forall a. a -> Maybe a
Just ([a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [a]
xs)

safeMinimum :: (Ord a) => [a] -> Maybe a
safeMinimum :: forall a. Ord a => [a] -> Maybe a
safeMinimum [] = Maybe a
forall a. Maybe a
Nothing
safeMinimum [a]
xs = a -> Maybe a
forall a. a -> Maybe a
Just ([a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [a]
xs)

safeInt8Mean :: Sequence -> Maybe Token
safeInt8Mean :: Keys -> Maybe Token
safeInt8Mean [] = Maybe Token
forall a. Maybe a
Nothing
safeInt8Mean Keys
xs = Token -> Maybe Token
forall a. a -> Maybe a
Just (Keys -> Token
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Keys
xs Token -> Token -> Token
forall a. Integral a => a -> a -> a
`div` Int -> Token
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Keys -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Keys
xs))

-- | Computes the "width", or number of nonzero entries, of the rows of a `Selector`.
selWidth :: Selector -> Sequence
selWidth :: Selector -> Keys
selWidth = ([Bool] -> Token) -> Selector -> Keys
forall a b. (a -> b) -> [a] -> [b]
map (Keys -> Token
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Keys -> Token) -> ([Bool] -> Keys) -> [Bool] -> Token
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool -> Token) -> [Bool] -> Keys
forall a b. (a -> b) -> [a] -> [b]
map Bool -> Token
fromBool)

fromBool :: Bool -> Token
fromBool :: Bool -> Token
fromBool Bool
True = Token
1
fromBool Bool
_ = Token
0

-- | Applies an elementwise operation to a sequence of tokens.
--
-- Roughly matches the MLP layer in a Transformer. Alias for `map`.
tokMap :: (Token -> Token) -> Sequence -> Sequence
tokMap :: (Token -> Token) -> Keys -> Keys
tokMap = (Token -> Token) -> Keys -> Keys
forall a b. (a -> b) -> [a] -> [b]
map

-- | Applies an elementwise operation for pairs of tokens on a pair of sequences.
-- Alias for `zipWith`.
seqMap :: (Token -> Token -> Token) -> Sequence -> Sequence -> Sequence
seqMap :: (Token -> Token -> Token) -> Keys -> Keys -> Keys
seqMap = (Token -> Token -> Token) -> Keys -> Keys -> Keys
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith

-- | Creates a sequence of the same length as the provided sequence filled with the provided token.
-- Alias for `filledWith`.
full :: Sequence -> Token -> Sequence
full :: Keys -> Token -> Keys
full = Keys -> Token -> Keys
filledWith

-- | Extracts the indices of the elements in a sequence.
-- Alias for `indicesOf`.
indices :: Sequence -> Sequence
indices :: Keys -> Keys
indices = Keys -> Keys
indicesOf

-- | Creates an aggregator with a given aggregation type.
-- Alias for `aggregate`.
aggr :: AggregationType -> Token -> Aggregator
aggr :: AggregationType -> Token -> Aggregator
aggr = AggregationType -> Token -> Aggregator
aggregate

-- | Produces a selector indicating which pairs of `Keys` and `Queries` match.
select ::
  -- | Whether to use causal selection
  Bool ->
  -- | A collection of `Keys` to check against `Queries`
  Keys ->
  -- | A collection of `Queries` to check against `Keys`
  Queries ->
  -- | A boolean predicate that determines whether a key and query match
  Predicate ->
  -- | A collection of boolean sequences indicating which pairs of `Keys` and `Queries` match
  Selector
select :: Bool -> Keys -> Keys -> Predicate -> Selector
select Bool
True = Keys -> Keys -> Predicate -> Selector
selectCausal
select Bool
False = Keys -> Keys -> Predicate -> Selector
selectAcausal

-- | Non-causal selection is included for some reason.
selectAcausal :: Keys -> Queries -> Predicate -> Selector
selectAcausal :: Keys -> Keys -> Predicate -> Selector
selectAcausal Keys
keys Keys
queries Predicate
predicate = [[Predicate
predicate Token
keyIndex Token
queryIndex | Token
keyIndex <- Keys
keys] | Token
queryIndex <- Keys
queries]