{-# LANGUAGE ScopedTypeVariables #-}

-- SPDX-License-Identifier: MPL-2.0

{- | CPU tensors with immutable boxed backing arrays and shared, zero-copy
strided views. Logical indexing is zero based; strides and offsets count
elements. Scalars have shape [], one element; zero axes have no elements.
No implicit numeric promotion occurs: 'promoteTensor' makes it explicit.
Garbage collection manages immutable backing storage and view lifetimes.
-}
module HMath.Tensor (
    Tensor,
    Layout (..),
    Device (..),
    tensor,
    tensorColumnMajor,
    tensorValues,
    tensorShape,
    tensorStrides,
    tensorOffset,
    tensorDType,
    tensorDevice,
    tensorLayout,
    isContiguous,
    contiguous,
    numel,
    shapeSize,
    inferBroadcast,
    tensorZeros,
    tensorOnes,
    tensorFull,
    tensorArange,
    tensorLinspace,
    tensorLogspace,
    tensorRandom,
    indexTensor,
    sliceTensor,
    fancyIndex,
    booleanIndex,
    gather,
    scatter,
    scatterAdd,
    reshapeTensor,
    flattenTensor,
    squeeze,
    unsqueeze,
    expand,
    permute,
    transposeTensor,
    moveAxis,
    swapAxes,
    broadcastTo,
    repeatTensor,
    tile,
    concatenate,
    stack,
    splitTensor,
    chunkTensor,
    padTensor,
    roll,
    flipTensor,
    rotate,
    sortTensor,
    argsort,
    topk,
    unique,
    nonzero,
    whereTensor,
    maskedFill,
    maskedSelect,
    mapTensor,
    zipTensor,
    tensorAdd,
    tensorSubtract,
    tensorMultiply,
    tensorDivide,
    tensorEqual,
    tensorLess,
    tensorAnd,
    tensorOr,
    tensorNot,
    promoteTensor,
    sumAxes,
    meanAxes,
    productAxes,
    minAxes,
    maxAxes,
    argmin,
    argmax,
    varianceAxes,
    normAxes,
    cumulativeSum,
    cumulativeProduct,
    tensordot,
    einsum,
    batchedMatmul,
    tensorKronecker,
    serializeTensor,
    deserializeTensor,
    TensorGraph,
    inputGraph,
    mapGraph,
    zipGraph,
    evaluateGraph,
) where

import Data.Array (Array, bounds, listArray, (!))
import Data.List (elemIndex, nub, sort, sortOn)
import qualified Data.Map.Strict as M
import Data.Proxy (Proxy (..))
import Data.Typeable (Typeable, typeRep)
import HMath.Error
import Prelude hiding (repeat)

{- | Tensor storage classification: canonical row major,column major,or general strided
view.
-}
data Layout = RowMajor | ColumnMajor | Strided deriving (Layout -> Layout -> Bool
(Layout -> Layout -> Bool)
-> (Layout -> Layout -> Bool) -> Eq Layout
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Layout -> Layout -> Bool
== :: Layout -> Layout -> Bool
$c/= :: Layout -> Layout -> Bool
/= :: Layout -> Layout -> Bool
Eq, Int -> Layout -> ShowS
[Layout] -> ShowS
Layout -> String
(Int -> Layout -> ShowS)
-> (Layout -> String) -> ([Layout] -> ShowS) -> Show Layout
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Layout -> ShowS
showsPrec :: Int -> Layout -> ShowS
$cshow :: Layout -> String
show :: Layout -> String
$cshowList :: [Layout] -> ShowS
showList :: [Layout] -> ShowS
Show, ReadPrec [Layout]
ReadPrec Layout
Int -> ReadS Layout
ReadS [Layout]
(Int -> ReadS Layout)
-> ReadS [Layout]
-> ReadPrec Layout
-> ReadPrec [Layout]
-> Read Layout
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS Layout
readsPrec :: Int -> ReadS Layout
$creadList :: ReadS [Layout]
readList :: ReadS [Layout]
$creadPrec :: ReadPrec Layout
readPrec :: ReadPrec Layout
$creadListPrec :: ReadPrec [Layout]
readListPrec :: ReadPrec [Layout]
Read)

-- | Tensor execution/storage device; this release supports CPU.
data Device = CPU deriving (Device -> Device -> Bool
(Device -> Device -> Bool)
-> (Device -> Device -> Bool) -> Eq Device
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Device -> Device -> Bool
== :: Device -> Device -> Bool
$c/= :: Device -> Device -> Bool
/= :: Device -> Device -> Bool
Eq, Int -> Device -> ShowS
[Device] -> ShowS
Device -> String
(Int -> Device -> ShowS)
-> (Device -> String) -> ([Device] -> ShowS) -> Show Device
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Device -> ShowS
showsPrec :: Int -> Device -> ShowS
$cshow :: Device -> String
show :: Device -> String
$cshowList :: [Device] -> ShowS
showList :: [Device] -> ShowS
Show, ReadPrec [Device]
ReadPrec Device
Int -> ReadS Device
ReadS [Device]
(Int -> ReadS Device)
-> ReadS [Device]
-> ReadPrec Device
-> ReadPrec [Device]
-> Read Device
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS Device
readsPrec :: Int -> ReadS Device
$creadList :: ReadS [Device]
readList :: ReadS [Device]
$creadPrec :: ReadPrec Device
readPrec :: ReadPrec Device
$creadListPrec :: ReadPrec [Device]
readListPrec :: ReadPrec [Device]
Read)

{- | Immutable boxed CPU tensor with signed element strides and shared view storage.
Construct through tensor and checked view operations; garbage collection manages
lifetime.
-}
data Tensor a = Tensor [Int] [Int] Int (Array Int a) deriving (Int -> Tensor a -> ShowS
[Tensor a] -> ShowS
Tensor a -> String
(Int -> Tensor a -> ShowS)
-> (Tensor a -> String) -> ([Tensor a] -> ShowS) -> Show (Tensor a)
forall a. Show a => Int -> Tensor a -> ShowS
forall a. Show a => [Tensor a] -> ShowS
forall a. Show a => Tensor a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Tensor a -> ShowS
showsPrec :: Int -> Tensor a -> ShowS
$cshow :: forall a. Show a => Tensor a -> String
show :: Tensor a -> String
$cshowList :: forall a. Show a => [Tensor a] -> ShowS
showList :: [Tensor a] -> ShowS
Show)

instance (Eq a) => Eq (Tensor a) where
    Tensor a
a == :: Tensor a -> Tensor a -> Bool
== Tensor a
b = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
a [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
b Bool -> Bool -> Bool
&& Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
a [a] -> [a] -> Bool
forall a. Eq a => a -> a -> Bool
== Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
b

{- | Validate nonnegative dimensions and return their product as Int. Scalar [] has one
element; a zero axis gives zero. Negative dimensions or an Int-overflowing product
return Left InvalidArgument. Haskell tensors do not inherit the C ABI's smaller
rank/element caps.
-}
shapeSize :: [Int] -> Either MathError Int
shapeSize :: [Int] -> Either MathError Int
shapeSize [Int]
dims
    | (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
dims = MathError -> Either MathError Int
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tensor: negative shape")
    | Integer
size Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) = MathError -> Either MathError Int
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tensor: element count overflow")
    | Bool
otherwise = Int -> Either MathError Int
forall a b. b -> Either a b
Right (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
size)
  where
    size :: Integer
size = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
dims)

stridesFor :: Layout -> [Int] -> [Int]
stridesFor :: Layout -> [Int] -> [Int]
stridesFor Layout
layout [Int]
dims
    | Int
0 Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
dims = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims) Int
0
    | Layout
layout Layout -> Layout -> Bool
forall a. Eq a => a -> a -> Bool
== Layout
ColumnMajor = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims) ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 [Int]
dims)
    | Bool
otherwise = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 [Int]
dims)

{- | Given shape then flat row-major values, construct an immutable CPU tensor. Require
exactly the shape product values; invalid dimensions/count return Left MathError.
Generic values are not validated for finiteness.
-}
tensor :: [Int] -> [a] -> Either MathError (Tensor a)
tensor :: forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int]
dims [a]
xs = do
    n <- [Int] -> Either MathError Int
shapeSize [Int]
dims
    if n /= length xs
        then Left (ElementCountMismatch n (length xs))
        else Right (Tensor dims (stridesFor RowMajor dims) 0 (listArray (0, n - 1) xs))

{- | Given shape then values in column-major storage order, construct a column-major CPU
tensor. Validate shape/count as tensor; logical tensorValues still traverses row-major
coordinates.
-}
tensorColumnMajor :: [Int] -> [a] -> Either MathError (Tensor a)
tensorColumnMajor :: forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensorColumnMajor [Int]
dims [a]
xs = do
    Tensor _ _ offset storage <- [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int]
dims [a]
xs
    Right (Tensor dims (stridesFor ColumnMajor dims) offset storage)

-- | Return dimensions in axis order. Scalar shape is []; empty axes are retained.
tensorShape :: Tensor a -> [Int]
tensorShape :: forall a. Tensor a -> [Int]
tensorShape (Tensor [Int]
dims [Int]
_ Int
_ Array Int a
_) = [Int]
dims

{- | Return signed element strides for each axis. These describe shared views; they are not
byte offsets and can be zero (broadcast) or negative (reverse slices).
-}
tensorStrides :: Tensor a -> [Int]
tensorStrides :: forall a. Tensor a -> [Int]
tensorStrides (Tensor [Int]
_ [Int]
strides Int
_ Array Int a
_) = [Int]
strides

{- | Return backing-array offset in elements of the logical coordinate origin. Combine with
tensorStrides for view indexing.
-}
tensorOffset :: Tensor a -> Int
tensorOffset :: forall a. Tensor a -> Int
tensorOffset (Tensor [Int]
_ [Int]
_ Int
offset Array Int a
_) = Int
offset

{- | Return the Typeable element type name as a String. This is diagnostic metadata, not an
implicit conversion or stable external dtype code.
-}
tensorDType :: forall a. (Typeable a) => Tensor a -> String
tensorDType :: forall a. Typeable a => Tensor a -> String
tensorDType Tensor a
_ = TypeRep -> String
forall a. Show a => a -> String
show (Proxy a -> TypeRep
forall {k} (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a))

-- | Return CPU, the supported device for this Haskell tensor implementation.
tensorDevice :: Tensor a -> Device
tensorDevice :: forall a. Tensor a -> Device
tensorDevice Tensor a
_ = Device
CPU

{- | Return logical element count (product of shape), including one for scalars and zero for
any empty axis.
-}
numel :: Tensor a -> Int
numel :: forall a. Tensor a -> Int
numel = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Int] -> Int) -> (Tensor a -> [Int]) -> Tensor a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape

{- | Test whether stored strides exactly match canonical row-major strides for this shape.
Offset does not prevent a contiguous view; column-major layouts generally return False.
-}
isContiguous :: Tensor a -> Bool
isContiguous :: forall a. Tensor a -> Bool
isContiguous Tensor a
t = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorStrides Tensor a
t [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== Layout -> [Int] -> [Int]
stridesFor Layout
RowMajor (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)

{- | Classify canonical row-major strides first, then column-major, otherwise Strided.
Ambiguous single-axis layouts are classified RowMajor.
-}
tensorLayout :: Tensor a -> Layout
tensorLayout :: forall a. Tensor a -> Layout
tensorLayout Tensor a
t
    | Tensor a -> Bool
forall a. Tensor a -> Bool
isContiguous Tensor a
t = Layout
RowMajor
    | Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorStrides Tensor a
t [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== Layout -> [Int] -> [Int]
stridesFor Layout
ColumnMajor (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) = Layout
ColumnMajor
    | Bool
otherwise = Layout
Strided

allIndices :: [Int] -> [[Int]]
allIndices :: [Int] -> [[Int]]
allIndices [Int]
dims | Int
0 Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
dims = []
allIndices [] = [[]]
allIndices (Int
d : [Int]
ds) = [Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
is | Int
i <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], [Int]
is <- [Int] -> [[Int]]
allIndices [Int]
ds]
linearIndex :: Tensor a -> [Int] -> Int
linearIndex :: forall a. Tensor a -> [Int] -> Int
linearIndex (Tensor [Int]
_ [Int]
strides Int
offset Array Int a
_) [Int]
ix = Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
strides [Int]
ix)
unsafeIndex :: Tensor a -> [Int] -> a
unsafeIndex :: forall a. Tensor a -> [Int] -> a
unsafeIndex t :: Tensor a
t@(Tensor [Int]
_ [Int]
_ Int
_ Array Int a
storage) [Int]
ix = Array Int a
storage Array Int a -> Int -> a
forall i e. Ix i => Array i e -> i -> e
! Tensor a -> [Int] -> Int
forall a. Tensor a -> [Int] -> Int
linearIndex Tensor a
t [Int]
ix

{- | Return logical values in row-major coordinate order, honoring strides/offsets rather
than exposing raw backing storage.
-}
tensorValues :: Tensor a -> [a]
tensorValues :: forall a. Tensor a -> [a]
tensorValues Tensor a
t = [Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex Tensor a
t [Int]
ix | [Int]
ix <- [Int] -> [[Int]]
allIndices (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)]

{- | Copy logical values into a new canonical row-major tensor of the same shape. Return
checked construction result; existing views remain unchanged.
-}
contiguous :: Tensor a -> Either MathError (Tensor a)
contiguous :: forall a. Tensor a -> Either MathError (Tensor a)
contiguous Tensor a
t = [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t)

validIndex :: [Int] -> [Int] -> Bool
validIndex :: [Int] -> [Int] -> Bool
validIndex [Int]
dims [Int]
ix = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ix Bool -> Bool -> Bool
&& [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i Int
d -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
d) [Int]
ix [Int]
dims)

{- | Given tensor then complete zero-based coordinate list, return the element. Wrong
coordinate rank or out-of-range index returns Left InvalidArgument.
-}
indexTensor :: Tensor a -> [Int] -> Either MathError a
indexTensor :: forall a. Tensor a -> [Int] -> Either MathError a
indexTensor Tensor a
t [Int]
ix
    | [Int] -> [Int] -> Bool
validIndex (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) [Int]
ix = a -> Either MathError a
forall a b. b -> Either a b
Right (Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex Tensor a
t [Int]
ix)
    | Bool
otherwise = MathError -> Either MathError a
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"indexTensor: index out of bounds")

view :: [Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
view :: forall a.
[Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
view [Int]
dims [Int]
strides Int
offset (Tensor [Int]
_ [Int]
_ Int
_ Array Int a
storage) = do
    n <- [Int] -> Either MathError Int
shapeSize [Int]
dims
    let (lo, hi) = bounds storage
        extremes = (Int -> Int -> Integer) -> [Int] -> [Int] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
d Int
s -> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
s) [Int]
dims [Int]
strides
        low = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
offset Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min Integer
0) [Integer]
extremes)
        high = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
offset Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
0) [Integer]
extremes)
    if length dims /= length strides
        || (n > 0 && (low < toInteger lo || high > toInteger hi))
        || (n == 0 && (offset < lo || toInteger offset > toInteger hi + 1))
        then Left (InvalidArgument "tensor view: storage bounds invalid")
        else Right (Tensor dims strides offset storage)
axisValid :: Int -> Tensor a -> Either MathError ()
axisValid :: forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    | Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) = MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tensor: axis out of range")
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()
replaceAt :: Int -> a -> [a] -> [a]
replaceAt :: forall a. Int -> a -> [a] -> [a]
replaceAt Int
i a
x [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs
insertAt :: Int -> a -> [a] -> [a]
insertAt :: forall a. Int -> a -> [a] -> [a]
insertAt Int
i a
x [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
i [a]
xs
removeAt :: Int -> [a] -> [a]
removeAt :: forall a. Int -> [a] -> [a]
removeAt Int
i [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs

{- | Given shape then fill value, construct an immutable row-major tensor with every element
equal to that value. Invalid/overflowing shape fails.
-}
tensorFull :: [Int] -> a -> Either MathError (Tensor a)
tensorFull :: forall a. [Int] -> a -> Either MathError (Tensor a)
tensorFull [Int]
dims a
x = do n <- [Int] -> Either MathError Int
shapeSize [Int]
dims; tensor dims (replicate n x)

{- | Construct a row-major tensor of zeros with the supplied shape. Invalid dimensions or
element-product overflow return Left MathError.
-}
tensorZeros :: (Num a) => [Int] -> Either MathError (Tensor a)
tensorZeros :: forall a. Num a => [Int] -> Either MathError (Tensor a)
tensorZeros [Int]
dims = [Int] -> a -> Either MathError (Tensor a)
forall a. [Int] -> a -> Either MathError (Tensor a)
tensorFull [Int]
dims a
0

{- | Construct a row-major tensor of ones with the supplied shape. Scalars contain one 1;
zero axes contain no values.
-}
tensorOnes :: (Num a) => [Int] -> Either MathError (Tensor a)
tensorOnes :: forall a. Num a => [Int] -> Either MathError (Tensor a)
tensorOnes [Int]
dims = [Int] -> a -> Either MathError (Tensor a)
forall a. [Int] -> a -> Either MathError (Tensor a)
tensorFull [Int]
dims a
1

{- | Given finite start, exclusive stop, and nonzero finite step, construct a rank-one
arithmetic progression. A step pointing away from stop yields empty output. Reject
invalid inputs or Int-size count overflow.
-}
tensorArange :: (RealFloat a) => a -> a -> a -> Either MathError (Tensor a)
tensorArange :: forall a. RealFloat a => a -> a -> a -> Either MathError (Tensor a)
tensorArange a
start a
stop a
step
    | (a -> Bool) -> [a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\a
x -> a -> Bool
forall a. RealFloat a => a -> Bool
isNaN a
x Bool -> Bool -> Bool
|| a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite a
x) [a
start, a
stop, a
step] Bool -> Bool -> Bool
|| a
step a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"arange: finite bounds and nonzero finite step required")
    | Integer
size Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int) = MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"arange: element count overflow")
    | Bool
otherwise =
        [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor
            [Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
size]
            [Rational -> a
forall a. Fractional a => Rational -> a
fromRational (a -> Rational
forall a. Real a => a -> Rational
toRational a
start Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
+ Integer -> Rational
forall a. Num a => Integer -> a
fromInteger Integer
i Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
* a -> Rational
forall a. Real a => a -> Rational
toRational a
step) | Integer
i <- [Integer
0 .. Integer
size Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]]
  where
    ratio :: Rational
ratio = (a -> Rational
forall a. Real a => a -> Rational
toRational a
stop Rational -> Rational -> Rational
forall a. Num a => a -> a -> a
- a -> Rational
forall a. Real a => a -> Rational
toRational a
start) Rational -> Rational -> Rational
forall a. Fractional a => a -> a -> a
/ a -> Rational
forall a. Real a => a -> Rational
toRational a
step
    size :: Integer
size = Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
0 (Rational -> Integer
forall b. Integral b => Rational -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling Rational
ratio)

{- | Given start, stop, then nonnegative count, return a rank-one evenly spaced sequence
including both endpoints when count>=2. Count zero gives empty, count one gives start.
Negative count fails; generic arithmetic is unchecked.
-}
tensorLinspace :: (Fractional a) => a -> a -> Int -> Either MathError (Tensor a)
tensorLinspace :: forall a.
Fractional a =>
a -> a -> Int -> Either MathError (Tensor a)
tensorLinspace a
start a
stop Int
n
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"linspace: negative count")
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int
0] []
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int
1] [a
start]
    | Bool
otherwise =
        [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor
            [Int
n]
            [ if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                then a
start
                else
                    if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                        then a
stop
                        else
                            a
start a -> a -> a
forall a. Num a => a -> a -> a
* (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
                                a -> a -> a
forall a. Num a => a -> a -> a
+ a
stop a -> a -> a
forall a. Num a => a -> a -> a
* (Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i a -> a -> a
forall a. Fractional a => a -> a -> a
/ Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
            | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
            ]

{- | Given base, starting exponent, ending exponent, then count, return base raised to
tensorLinspace exponents. Negative count fails; Floating domain/overflow behavior
follows the element type.
-}
tensorLogspace :: (Floating a) => a -> a -> a -> Int -> Either MathError (Tensor a)
tensorLogspace :: forall a.
Floating a =>
a -> a -> a -> Int -> Either MathError (Tensor a)
tensorLogspace a
base a
start a
stop Int
n = (a -> a) -> Tensor a -> Tensor a
forall a b. (a -> b) -> Tensor a -> Tensor b
mapTensor (a
base a -> a -> a
forall a. Floating a => a -> a -> a
**) (Tensor a -> Tensor a)
-> Either MathError (Tensor a) -> Either MathError (Tensor a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> a -> Int -> Either MathError (Tensor a)
forall a.
Fractional a =>
a -> a -> Int -> Either MathError (Tensor a)
tensorLinspace a
start a
stop Int
n

{- | Explicit deterministic LCG seed. Values are uniform in [0,1); the returned
seed continues the stream, and does not access global mutable generator state.
-}
tensorRandom :: Integer -> [Int] -> Either MathError (Tensor Double, Integer)
tensorRandom :: Integer -> [Int] -> Either MathError (Tensor Double, Integer)
tensorRandom Integer
seed [Int]
dims = do
    n <- [Int] -> Either MathError Int
shapeSize [Int]
dims
    let states =
            Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take
                Int
n
                (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
drop Int
1 ((Integer -> Integer) -> Integer -> [Integer]
forall a. (a -> a) -> a -> [a]
iterate (\Integer
x -> (Integer
1664525 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1013904223) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4294967296) (Integer
seed Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4294967296)))
    t <- tensor dims [fromInteger x / 4294967296 | x <- states]
    Right (t, case reverse states of Integer
x : [Integer]
_ -> Integer
x; [] -> Integer
seed Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Integer
4294967296)

{- | Given axis,start,count,nonzero step,then tensor, return a zero-copy strided view.
Negative steps are supported with valid reachable coordinates. Invalid axis/extent,
stride overflow or backing bounds return Left MathError.
-}
sliceTensor :: Int -> Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
sliceTensor :: forall a.
Int -> Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
sliceTensor Int
axis Int
start Int
count Int
step Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    let dims :: [Int]
dims = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t
        strides :: [Int]
strides = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorStrides Tensor a
t
        d :: Int
d = [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
        end :: Integer
end = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
start Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
step
        offset :: Integer
offset = Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Tensor a -> Int
forall a. Tensor a -> Int
tensorOffset Tensor a
t) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
start Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
strides [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)
        stride :: Integer
stride = Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
strides [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
step
    if Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        Bool -> Bool -> Bool
|| Int
step Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        Bool -> Bool -> Bool
|| Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        Bool -> Bool -> Bool
|| (Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& (Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
d Bool -> Bool -> Bool
|| Integer
end Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
end Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
d))
        Bool -> Bool -> Bool
|| (Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
d)
        Bool -> Bool -> Bool
|| Integer
stride Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
minBound :: Int)
        Bool -> Bool -> Bool
|| Integer
stride Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
        then MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sliceTensor: invalid extent or step")
        else
            [Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
[Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
view
                (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis Int
count [Int]
dims)
                (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
stride) [Int]
strides)
                (if Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Tensor a -> Int
forall a. Tensor a -> Int
tensorOffset Tensor a
t else Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
offset)
                Tensor a
t

{- | Given axis then index list then tensor, copy selected positions along that axis in
index-list order, retaining duplicates. Invalid axis/indices or output-size overflow
fail.
-}
fancyIndex :: Int -> [Int] -> Tensor a -> Either MathError (Tensor a)
fancyIndex :: forall a. Int -> [Int] -> Tensor a -> Either MathError (Tensor a)
fancyIndex Int
axis [Int]
indices Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    let dims :: [Int]
dims = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t
    if (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) [Int]
indices
        then MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"fancyIndex: index out of bounds")
        else do
            let output :: [Int]
output = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices) [Int]
dims
            _ <- [Int] -> Either MathError Int
shapeSize [Int]
output
            tensor
                output
                [unsafeIndex t (replaceAt axis (indices !! (ix !! axis)) ix) | ix <- allIndices output]

{- | Given an identically shaped Bool mask then tensor, return a rank-one tensor of values
whose mask is True, in logical row-major order. Mask mismatch fails; alias of
maskedSelect.
-}
booleanIndex :: Tensor Bool -> Tensor a -> Either MathError (Tensor a)
booleanIndex :: forall a. Tensor Bool -> Tensor a -> Either MathError (Tensor a)
booleanIndex = Tensor Bool -> Tensor a -> Either MathError (Tensor a)
forall a. Tensor Bool -> Tensor a -> Either MathError (Tensor a)
maskedSelect

{- | Alias of fancyIndex: given axis, zero-based index list, then tensor, copy those axis
positions. Invalid axis/indices fail; index-list order and duplicates are preserved.
-}
gather :: Int -> [Int] -> Tensor a -> Either MathError (Tensor a)
gather :: forall a. Int -> [Int] -> Tensor a -> Either MathError (Tensor a)
gather = Int -> [Int] -> Tensor a -> Either MathError (Tensor a)
forall a. Int -> [Int] -> Tensor a -> Either MathError (Tensor a)
fancyIndex

scatterWith :: (a -> a -> a) -> Int -> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
scatterWith :: forall a.
(a -> a -> a)
-> Int
-> [Int]
-> Tensor a
-> Tensor a
-> Either MathError (Tensor a)
scatterWith a -> a -> a
combine Int
axis [Int]
indices Tensor a
updates Tensor a
target = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
target
    let dims :: [Int]
dims = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
target; updateDims :: [Int]
updateDims = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
indices) [Int]
dims
    if Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
updates [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int]
updateDims Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) [Int]
indices
        then MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"scatter: update shape or indices invalid")
        else do
            let initial :: Map [Int] a
initial = [([Int], a)] -> Map [Int] a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([[Int]] -> [a] -> [([Int], a)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [[Int]]
allIndices [Int]
dims) (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
target))
                apply :: Map [Int] a -> [Int] -> Map [Int] a
apply Map [Int] a
values [Int]
ix =
                    (a -> a) -> [Int] -> Map [Int] a -> Map [Int] a
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust
                        (\a
old -> a -> a -> a
combine a
old (Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex Tensor a
updates [Int]
ix))
                        (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis ([Int]
indices [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)) [Int]
ix)
                        Map [Int] a
values
                final :: Map [Int] a
final = (Map [Int] a -> [Int] -> Map [Int] a)
-> Map [Int] a -> [[Int]] -> Map [Int] a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Map [Int] a -> [Int] -> Map [Int] a
apply Map [Int] a
initial ([Int] -> [[Int]]
allIndices [Int]
updateDims)
            [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int]
dims [Map [Int] a
final Map [Int] a -> [Int] -> a
forall k a. Ord k => Map k a -> k -> a
M.! [Int]
ix | [Int]
ix <- [Int] -> [[Int]]
allIndices [Int]
dims]

{- | Given axis,index list,update tensor,then target, return a new target-shaped tensor with
updates written to selected positions. Update shape replaces that axis with index-list
length; invalid shape/indices fail. Repeated indices use the last update in traversal
order.
-}
scatter :: Int -> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
scatter :: forall a.
Int -> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
scatter = (a -> a -> a)
-> Int
-> [Int]
-> Tensor a
-> Tensor a
-> Either MathError (Tensor a)
forall a.
(a -> a -> a)
-> Int
-> [Int]
-> Tensor a
-> Tensor a
-> Either MathError (Tensor a)
scatterWith (\a
_ a
new -> a
new)

{- | Given axis,indices,updates,target, return a new target with updates added at selected
coordinates. Duplicate indices accumulate; update shape and indices follow scatter
validation.
-}
scatterAdd :: (Num a) => Int -> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
scatterAdd :: forall a.
Num a =>
Int -> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
scatterAdd = (a -> a -> a)
-> Int
-> [Int]
-> Tensor a
-> Tensor a
-> Either MathError (Tensor a)
forall a.
(a -> a -> a)
-> Int
-> [Int]
-> Tensor a
-> Tensor a
-> Either MathError (Tensor a)
scatterWith a -> a -> a
forall a. Num a => a -> a -> a
(+)

{- | Given new dimensions then tensor, preserve logical element order with unchanged count.
Contiguous input yields a shared view; non-contiguous input is copied. Invalid
dimensions/count fail; no -1 inference.
-}
reshapeTensor :: [Int] -> Tensor a -> Either MathError (Tensor a)
reshapeTensor :: forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
reshapeTensor [Int]
dims Tensor a
t = do
    size <- [Int] -> Either MathError Int
shapeSize [Int]
dims
    if size /= numel t
        then Left (ElementCountMismatch size (numel t))
        else
            if isContiguous t
                then view dims (stridesFor RowMajor dims) (tensorOffset t) t
                else tensor dims (tensorValues t)

{- | Reshape to rank one with numel entries in logical row-major order. Shares contiguous
storage and copies non-contiguous input; checked reshape errors propagate.
-}
flattenTensor :: Tensor a -> Either MathError (Tensor a)
flattenTensor :: forall a. Tensor a -> Either MathError (Tensor a)
flattenTensor Tensor a
t = [Int] -> Tensor a -> Either MathError (Tensor a)
forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
reshapeTensor [Tensor a -> Int
forall a. Tensor a -> Int
numel Tensor a
t] Tensor a
t

{- | Remove every length-one axis as a shared view, preserving storage and logical values. A
single-element all-one shape becomes scalar [].
-}
squeeze :: Tensor a -> Either MathError (Tensor a)
squeeze :: forall a. Tensor a -> Either MathError (Tensor a)
squeeze Tensor a
t =
    [Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
[Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
view
        [Int
d | Int
d <- Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t, Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1]
        [Int
s | (Int
d, Int
s) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorStrides Tensor a
t), Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1]
        (Tensor a -> Int
forall a. Tensor a -> Int
tensorOffset Tensor a
t)
        Tensor a
t

{- | Given insertion axis in [0,rank] then tensor, insert a length-one dimension as a shared
view. Out-of-range axis returns Left InvalidArgument.
-}
unsqueeze :: Int -> Tensor a -> Either MathError (Tensor a)
unsqueeze :: forall a. Int -> Tensor a -> Either MathError (Tensor a)
unsqueeze Int
axis Tensor a
t
    | Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) = MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"unsqueeze: axis out of range")
    | Bool
otherwise =
        [Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
[Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
view (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
insertAt Int
axis Int
1 (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)) (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
insertAt Int
axis Int
0 (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorStrides Tensor a
t)) (Tensor a -> Int
forall a. Tensor a -> Int
tensorOffset Tensor a
t) Tensor a
t

{- | Given a permutation of all zero-based axes then tensor, return a zero-copy view with
shape/strides reordered. Missing, duplicate or invalid axes fail.
-}
permute :: [Int] -> Tensor a -> Either MathError (Tensor a)
permute :: forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
permute [Int]
order Tensor a
t
    | [Int] -> [Int]
forall a. Ord a => [a] -> [a]
sort [Int]
order [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"permute: axes must be a permutation")
    | Bool
otherwise =
        [Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
[Int] -> [Int] -> Int -> Tensor a -> Either MathError (Tensor a)
view ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
order) ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorStrides Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
order) (Tensor a -> Int
forall a. Tensor a -> Int
tensorOffset Tensor a
t) Tensor a
t

{- | Reverse every axis of a tensor as a shared view. For rank two this exchanges
rows/columns; it is not restricted to matrices.
-}
transposeTensor :: Tensor a -> Either MathError (Tensor a)
transposeTensor :: forall a. Tensor a -> Either MathError (Tensor a)
transposeTensor Tensor a
t = [Int] -> Tensor a -> Either MathError (Tensor a)
forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
permute ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]) Tensor a
t

{- | Given source axis,destination axis,then tensor, move one axis while preserving the
order of others. Both axes must identify existing dimensions; invalid axes fail.
-}
moveAxis :: Int -> Int -> Tensor a -> Either MathError (Tensor a)
moveAxis :: forall a. Int -> Int -> Tensor a -> Either MathError (Tensor a)
moveAxis Int
source Int
destination Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
source Tensor a
t
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
destination Tensor a
t
    [Int] -> Tensor a -> Either MathError (Tensor a)
forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
permute (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
insertAt Int
destination Int
source ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
source) [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])) Tensor a
t

{- | Given two existing axes then tensor, exchange their dimensions/strides in a shared
view. Invalid axes return Left MathError.
-}
swapAxes :: Int -> Int -> Tensor a -> Either MathError (Tensor a)
swapAxes :: forall a. Int -> Int -> Tensor a -> Either MathError (Tensor a)
swapAxes Int
a Int
b Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
a Tensor a
t
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
b Tensor a
t
    [Int] -> Tensor a -> Either MathError (Tensor a)
forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
permute [if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
a then Int
b else if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
b then Int
a else Int
i | Int
i <- [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Tensor a
t

{- | Given two shapes, align trailing dimensions; equal sizes or a size-one dimension are
compatible. Return the broadcasted shape after nonnegative/Int-size validation;
incompatible dimensions fail.
-}
inferBroadcast :: [Int] -> [Int] -> Either MathError [Int]
inferBroadcast :: [Int] -> [Int] -> Either MathError [Int]
inferBroadcast [Int]
as [Int]
bs = do
    _ <- [Int] -> Either MathError Int
shapeSize [Int]
as
    _ <- shapeSize bs
    dims <- sequence (zipWith compatible (pad as) (pad bs))
    _ <- shapeSize dims
    Right dims
  where
    n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
as) ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
bs)
    pad :: [a] -> [a]
pad [a]
xs = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) a
1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
xs
    compatible :: a -> a -> Either MathError a
compatible a
a a
b
        | a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
b = a -> Either MathError a
forall a b. b -> Either a b
Right a
a
        | a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 = a -> Either MathError a
forall a b. b -> Either a b
Right a
b
        | a
b a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 = a -> Either MathError a
forall a b. b -> Either a b
Right a
a
        | Bool
otherwise = MathError -> Either MathError a
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"broadcast: incompatible shapes")

{- | Given target shape then tensor, return a shared zero-stride view where
existing/trailing size-one axes can expand. Target must be exactly a valid broadcast
destination; incompatible/overflowing shapes fail.
-}
broadcastTo :: [Int] -> Tensor a -> Either MathError (Tensor a)
broadcastTo :: forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
broadcastTo [Int]
dims Tensor a
t = do
    inferred <- [Int] -> [Int] -> Either MathError [Int]
inferBroadcast [Int]
dims (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)
    if inferred /= dims || length dims < length (tensorShape t)
        then Left (InvalidArgument "broadcastTo: incompatible target")
        else do
            let rank = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims
                missing = Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)
                source = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
missing Int
1 [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t
                strides = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
missing Int
0 [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorStrides Tensor a
t
            view
                dims
                (zipWith3 (\Int
d Int
old Int
s -> if Int
old Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 Bool -> Bool -> Bool
&& Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1 then Int
0 else Int
s) dims source strides)
                (tensorOffset t)
                t

{- | Alias of broadcastTo: expand singleton dimensions to a requested compatible target
shape without copying values. Invalid broadcast shapes fail.
-}
expand :: [Int] -> Tensor a -> Either MathError (Tensor a)
expand :: forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
expand = [Int] -> Tensor a -> Either MathError (Tensor a)
forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
broadcastTo

{- | Given axis,repetition count,then tensor, repeat each individual axis position
consecutively. Nonnegative count is required; zero yields an empty axis. Invalid
axis/count or size overflow fails.
-}
repeatTensor :: Int -> Int -> Tensor a -> Either MathError (Tensor a)
repeatTensor :: forall a. Int -> Int -> Tensor a -> Either MathError (Tensor a)
repeatTensor Int
axis Int
repetitions Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    let dim :: Int
dim = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
        new :: Integer
new = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
dim Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
repetitions
    if Int
repetitions Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Integer
new Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
        then MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"repeatTensor: invalid repetition or size overflow")
        else do
            let dims :: [Int]
dims = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
new) (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)
            _ <- [Int] -> Either MathError Int
shapeSize [Int]
dims
            tensor
                dims
                [unsafeIndex t (replaceAt axis (ix !! axis `div` repetitions) ix) | ix <- allIndices dims]

{- | Given one nonnegative repetition count per existing axis then tensor, tile whole axis
contents into repeated blocks. Rank mismatch, negative repetitions or resulting size
overflow fail.
-}
tile :: [Int] -> Tensor a -> Either MathError (Tensor a)
tile :: forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
tile [Int]
repetitions Tensor a
t
    | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
repetitions Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
repetitions =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tile: repetition rank mismatch or negative repetition")
    | (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)) [Integer]
dimensions = MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tile: shape overflow")
    | Bool
otherwise = do
        let dims :: [Int]
dims = (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger [Integer]
dimensions
        _ <- [Int] -> Either MathError Int
shapeSize [Int]
dims
        tensor dims [unsafeIndex t (zipWith mod ix (tensorShape t)) | ix <- allIndices dims]
  where
    dimensions :: [Integer]
dimensions = (Int -> Int -> Integer) -> [Int] -> [Int] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
d Int
n -> Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n) (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) [Int]
repetitions

{- | Given existing axis then nonempty tensor list, copy tensors end-to-end along that axis.
All other dimensions must match; empty list, axis/shape mismatch or size overflow
fails.
-}
concatenate :: Int -> [Tensor a] -> Either MathError (Tensor a)
concatenate :: forall a. Int -> [Tensor a] -> Either MathError (Tensor a)
concatenate Int
_ [] = MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"concatenate: nonempty input required")
concatenate Int
axis ts :: [Tensor a]
ts@(Tensor a
first : [Tensor a]
_) = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
first
    let dims :: [Int]
dims = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
first
        total :: Integer
total = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) | Tensor a
t <- [Tensor a]
ts, [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims]
    if (Tensor a -> Bool) -> [Tensor a] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any
        (\Tensor a
t -> [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Bool -> Bool -> Bool
|| Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
removeAt Int
axis (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
removeAt Int
axis [Int]
dims)
        [Tensor a]
ts
        Bool -> Bool -> Bool
|| Integer
total Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)
        then MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"concatenate: incompatible shapes or dimension overflow")
        else do
            let output :: [Int]
output = Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
total) [Int]
dims
                locate :: Int -> [Tensor a] -> (Tensor a, Int)
locate Int
position (Tensor a
t : [Tensor a]
rest)
                    | Int
position Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis = (Tensor a
t, Int
position)
                    | Bool
otherwise = Int -> [Tensor a] -> (Tensor a, Int)
locate (Int
position Int -> Int -> Int
forall a. Num a => a -> a -> a
- Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) [Tensor a]
rest
                locate Int
_ [] = String -> (Tensor a, Int)
forall a. HasCallStack => String -> a
error String
"concatenate: internal index invariant"
            _ <- [Int] -> Either MathError Int
shapeSize [Int]
output
            tensor
                output
                [ let (t, j) = locate (ix !! axis) ts in unsafeIndex t (replaceAt axis j ix) | ix <- allIndices output
                ]

{- | Given insertion axis then nonempty list of equally shaped tensors, insert that axis and
concatenate along it. Invalid insertion axis or incompatible input shapes fail.
-}
stack :: Int -> [Tensor a] -> Either MathError (Tensor a)
stack :: forall a. Int -> [Tensor a] -> Either MathError (Tensor a)
stack Int
axis [Tensor a]
ts = (Tensor a -> Either MathError (Tensor a))
-> [Tensor a] -> Either MathError [Tensor a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (Int -> Tensor a -> Either MathError (Tensor a)
forall a. Int -> Tensor a -> Either MathError (Tensor a)
unsqueeze Int
axis) [Tensor a]
ts Either MathError [Tensor a]
-> ([Tensor a] -> Either MathError (Tensor a))
-> Either MathError (Tensor a)
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> [Tensor a] -> Either MathError (Tensor a)
forall a. Int -> [Tensor a] -> Either MathError (Tensor a)
concatenate Int
axis

{- | Given axis,nonnegative segment sizes,then tensor, return shared slices partitioning
that axis exactly. Sizes must sum to the axis length; invalid axis/partition fails.
-}
splitTensor :: Int -> [Int] -> Tensor a -> Either MathError [Tensor a]
splitTensor :: forall a. Int -> [Int] -> Tensor a -> Either MathError [Tensor a]
splitTensor Int
axis [Int]
sizes Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    if (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [Int]
sizes Bool -> Bool -> Bool
|| [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
sizes) Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)
        then MathError -> Either MathError [Tensor a]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"splitTensor: sizes must partition axis")
        else
            [Either MathError (Tensor a)] -> Either MathError [Tensor a]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [Int -> Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
Int -> Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
sliceTensor Int
axis Int
start Int
count Int
1 Tensor a
t | (Int
start, Int
count) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 [Int]
sizes) [Int]
sizes]

{- | Given axis,positive chunk count,then tensor, split into exactly that many near-equal
shared slices; earlier chunks get the remainder. Empty chunks are allowed if chunks
exceed axis size. Invalid axis/count fails.
-}
chunkTensor :: Int -> Int -> Tensor a -> Either MathError [Tensor a]
chunkTensor :: forall a. Int -> Int -> Tensor a -> Either MathError [Tensor a]
chunkTensor Int
axis Int
chunks Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    if Int
chunks Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
        then MathError -> Either MathError [Tensor a]
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"chunkTensor: positive chunk count required")
        else do
            let n :: Int
n = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis; (Int
q, Int
r) = Int
n Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
chunks
            Int -> [Int] -> Tensor a -> Either MathError [Tensor a]
forall a. Int -> [Int] -> Tensor a -> Either MathError [Tensor a]
splitTensor Int
axis [Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r then Int
1 else Int
0 | Int
i <- [Int
0 .. Int
chunks Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]] Tensor a
t

{- | Given (before,after) nonnegative padding for every axis,fill value,then tensor, copy
into a padded output. Rank mismatch, negative padding or output-size overflow fails.
-}
padTensor :: [(Int, Int)] -> a -> Tensor a -> Either MathError (Tensor a)
padTensor :: forall a.
[(Int, Int)] -> a -> Tensor a -> Either MathError (Tensor a)
padTensor [(Int, Int)]
padding a
value Tensor a
t
    | [(Int, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Int)]
padding Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) Bool -> Bool -> Bool
|| ((Int, Int) -> Bool) -> [(Int, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
a, Int
b) -> Int
a Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) [(Int, Int)]
padding =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"padTensor: invalid padding")
    | (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)) [Integer]
dimensions =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"padTensor: dimension overflow")
    | Bool
otherwise = do
        let output :: [Int]
output = (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger [Integer]
dimensions
        _ <- [Int] -> Either MathError Int
shapeSize [Int]
output
        tensor
            output
            [ let old = (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (-) [Int]
ix (((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> a
fst [(Int, Int)]
padding)
               in if validIndex (tensorShape t) old then unsafeIndex t old else value
            | ix <- allIndices output
            ]
  where
    dimensions :: [Integer]
dimensions = (Int -> (Int, Int) -> Integer)
-> [Int] -> [(Int, Int)] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
d (Int
a, Int
b) -> Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
b) (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) [(Int, Int)]
padding

{- | Given axis,signed shift,then tensor, cyclically shift elements along the axis,
preserving shape. Positive shifts move values toward larger indices. Invalid axis
fails; the output is copied.
-}
roll :: Int -> Int -> Tensor a -> Either MathError (Tensor a)
roll :: forall a. Int -> Int -> Tensor a -> Either MathError (Tensor a)
roll Int
axis Int
shift Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    let dims :: [Int]
dims = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t; n :: Int
n = [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
    [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor
        [Int]
dims
        [ Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex
            Tensor a
t
            (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis (Integer -> Int
forall a. Num a => Integer -> a
fromInteger ((Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
shift) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`mod` Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
n)) [Int]
ix)
        | [Int]
ix <- [Int] -> [[Int]]
allIndices [Int]
dims
        ]

{- | Given axis then tensor, reverse that axis using a negative-stride shared slice. Invalid
axis returns Left MathError.
-}
flipTensor :: Int -> Tensor a -> Either MathError (Tensor a)
flipTensor :: forall a. Int -> Tensor a -> Either MathError (Tensor a)
flipTensor Int
axis Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    let n :: Int
n = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis
    Int -> Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
Int -> Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
sliceTensor Int
axis (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Int
n (-Int
1) Tensor a
t

{- | Given two distinct axes,integer quarter-turn count,then tensor, repeatedly swap axes
and flip the first. Turns are reduced modulo four; invalid/equal axes fail.
-}
rotate :: Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
rotate :: forall a.
Int -> Int -> Int -> Tensor a -> Either MathError (Tensor a)
rotate Int
axisA Int
axisB Int
turns Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axisA Tensor a
t
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axisB Tensor a
t
    if Int
axisA Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
axisB
        then MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"rotate: distinct axes required")
        else Int -> Tensor a -> Either MathError (Tensor a)
forall {t} {a}.
(Eq t, Num t) =>
t -> Tensor a -> Either MathError (Tensor a)
go (Int
turns Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
4) Tensor a
t
  where
    go :: t -> Tensor a -> Either MathError (Tensor a)
go t
0 Tensor a
x = Tensor a -> Either MathError (Tensor a)
forall a b. b -> Either a b
Right Tensor a
x
    go t
k Tensor a
x = Int -> Int -> Tensor a -> Either MathError (Tensor a)
forall a. Int -> Int -> Tensor a -> Either MathError (Tensor a)
swapAxes Int
axisA Int
axisB Tensor a
x Either MathError (Tensor a)
-> (Tensor a -> Either MathError (Tensor a))
-> Either MathError (Tensor a)
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Tensor a -> Either MathError (Tensor a)
forall a. Int -> Tensor a -> Either MathError (Tensor a)
flipTensor Int
axisA Either MathError (Tensor a)
-> (Tensor a -> Either MathError (Tensor a))
-> Either MathError (Tensor a)
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= t -> Tensor a -> Either MathError (Tensor a)
go (t
k t -> t -> t
forall a. Num a => a -> a -> a
- t
1)

{- | Flatten logical values and return a rank-one tensor sorted ascending under Ord.
Original shape is not preserved.
-}
sortTensor :: (Ord a) => Tensor a -> Either MathError (Tensor a)
sortTensor :: forall a. Ord a => Tensor a -> Either MathError (Tensor a)
sortTensor Tensor a
t = [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Tensor a -> Int
forall a. Tensor a -> Int
numel Tensor a
t] ([a] -> [a]
forall a. Ord a => [a] -> [a]
sort (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t))

{- | Flatten logical values and return a rank-one tensor of stable zero-based ascending-sort
indices. Ties preserve original order.
-}
argsort :: (Ord a) => Tensor a -> Either MathError (Tensor Int)
argsort :: forall a. Ord a => Tensor a -> Either MathError (Tensor Int)
argsort Tensor a
t = [Int] -> [Int] -> Either MathError (Tensor Int)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Tensor a -> Int
forall a. Tensor a -> Int
numel Tensor a
t] (((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst (((Int, a) -> a) -> [(Int, a)] -> [(Int, a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int, a) -> a
forall a b. (a, b) -> b
snd ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t))))

{- | Given k in [0,numel] then tensor, return rank-one (values,flat indices) for k largest
values in descending order. Equal values use reverse stable-sort order. Invalid k
fails.
-}
topk :: (Ord a) => Int -> Tensor a -> Either MathError (Tensor a, Tensor Int)
topk :: forall a.
Ord a =>
Int -> Tensor a -> Either MathError (Tensor a, Tensor Int)
topk Int
k Tensor a
t
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Tensor a -> Int
forall a. Tensor a -> Int
numel Tensor a
t = MathError -> Either MathError (Tensor a, Tensor Int)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"topk: invalid k")
    | Bool
otherwise =
        let selected :: [(Int, a)]
selected = Int -> [(Int, a)] -> [(Int, a)]
forall a. Int -> [a] -> [a]
take Int
k ([(Int, a)] -> [(Int, a)]
forall a. [a] -> [a]
reverse (((Int, a) -> a) -> [(Int, a)] -> [(Int, a)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Int, a) -> a
forall a b. (a, b) -> b
snd ([Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t))))
         in (,) (Tensor a -> Tensor Int -> (Tensor a, Tensor Int))
-> Either MathError (Tensor a)
-> Either MathError (Tensor Int -> (Tensor a, Tensor Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int
k] (((Int, a) -> a) -> [(Int, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> a
forall a b. (a, b) -> b
snd [(Int, a)]
selected) Either MathError (Tensor Int -> (Tensor a, Tensor Int))
-> Either MathError (Tensor Int)
-> Either MathError (Tensor a, Tensor Int)
forall a b.
Either MathError (a -> b)
-> Either MathError a -> Either MathError b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Int] -> [Int] -> Either MathError (Tensor Int)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int
k] (((Int, a) -> Int) -> [(Int, a)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> Int
forall a b. (a, b) -> a
fst [(Int, a)]
selected)

{- | Return sorted distinct logical values as a rank-one tensor. Input shape/order are not
preserved.
-}
unique :: (Ord a) => Tensor a -> Either MathError (Tensor a)
unique :: forall a. Ord a => Tensor a -> Either MathError (Tensor a)
unique Tensor a
t = let xs :: [a]
xs = Map a () -> [a]
forall k a. Map k a -> [k]
M.keys ([(a, ())] -> Map a ()
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(a
x, ()) | a
x <- Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t]) in [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [[a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs] [a]
xs

{- | Return a count-by-rank tensor of zero-based coordinates of nonzero elements in
row-major traversal order. For a nonzero scalar the shape is [1,0].
-}
nonzero :: (Eq a, Num a) => Tensor a -> Either MathError (Tensor Int)
nonzero :: forall a.
(Eq a, Num a) =>
Tensor a -> Either MathError (Tensor Int)
nonzero Tensor a
t =
    let ix :: [[Int]]
ix = [[Int]
is | [Int]
is <- [Int] -> [[Int]]
allIndices (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t), Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex Tensor a
t [Int]
is a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0]
     in [Int] -> [Int] -> Either MathError (Tensor Int)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [[[Int]] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
ix, [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)] ([[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Int]]
ix)

{- | Given Bool mask,yes tensor,no tensor, broadcast all three shapes and select yes/no
values elementwise. Incompatible/overflowing broadcast shapes return Left MathError.
-}
whereTensor :: Tensor Bool -> Tensor a -> Tensor a -> Either MathError (Tensor a)
whereTensor :: forall a.
Tensor Bool -> Tensor a -> Tensor a -> Either MathError (Tensor a)
whereTensor Tensor Bool
mask Tensor a
yes Tensor a
no = do
    dims <- [Int] -> [Int] -> Either MathError [Int]
inferBroadcast (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
yes) (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
no) Either MathError [Int]
-> ([Int] -> Either MathError [Int]) -> Either MathError [Int]
forall a b.
Either MathError a
-> (a -> Either MathError b) -> Either MathError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= [Int] -> [Int] -> Either MathError [Int]
inferBroadcast (Tensor Bool -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor Bool
mask)
    m <- broadcastTo dims mask
    y <- broadcastTo dims yes
    n <- broadcastTo dims no
    tensor
        dims
        ( zipWith3
            (\Bool
condition a
a a
b -> if Bool
condition then a
a else a
b)
            (tensorValues m)
            (tensorValues y)
            (tensorValues n)
        )

{- | Given broadcastable mask,fill value,then tensor, return a copied result with True
positions replaced. Incompatible mask broadcast or output construction fails.
-}
maskedFill :: Tensor Bool -> a -> Tensor a -> Either MathError (Tensor a)
maskedFill :: forall a.
Tensor Bool -> a -> Tensor a -> Either MathError (Tensor a)
maskedFill Tensor Bool
mask a
value Tensor a
t = do full <- [Int] -> a -> Either MathError (Tensor a)
forall a. [Int] -> a -> Either MathError (Tensor a)
tensorFull (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) a
value; whereTensor mask full t

{- | Given mask then identically shaped tensor, return a rank-one copy of True-selected
values in row-major order. Shape mismatch returns Left InvalidArgument.
-}
maskedSelect :: Tensor Bool -> Tensor a -> Either MathError (Tensor a)
maskedSelect :: forall a. Tensor Bool -> Tensor a -> Either MathError (Tensor a)
maskedSelect Tensor Bool
mask Tensor a
t
    | Tensor Bool -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor Bool
mask [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
/= Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t = MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"maskedSelect: mask shape mismatch")
    | Bool
otherwise =
        let xs :: [a]
xs = [a
x | (Bool
True, a
x) <- [Bool] -> [a] -> [(Bool, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Tensor Bool -> [Bool]
forall a. Tensor a -> [a]
tensorValues Tensor Bool
mask) (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t)] in [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [[a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs] [a]
xs

{- | Apply scalar arithmetic, comparisons, logic, or elementary functions such
as @sin@, @exp@, and @log@ to each logical element. A contiguous output is
created; scalar numeric domain/overflow behavior follows the supplied function.
-}
mapTensor :: (a -> b) -> Tensor a -> Tensor b
mapTensor :: forall a b. (a -> b) -> Tensor a -> Tensor b
mapTensor a -> b
f Tensor a
t = case [Int] -> [b] -> Either MathError (Tensor b)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) ((a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t)) of
    Right Tensor b
r -> Tensor b
r
    Left MathError
_ -> String -> Tensor b
forall a. HasCallStack => String -> a
error String
"mapTensor: shape invariant"

{- | Given binary element function then two tensors, broadcast compatible shapes and return
a new contiguous result. Shape/size failures return Left MathError; function
evaluation/domain errors follow the supplied function.
-}
zipTensor :: (a -> b -> c) -> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor :: forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor a -> b -> c
f Tensor a
a Tensor b
b = do
    dims <- [Int] -> [Int] -> Either MathError [Int]
inferBroadcast (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
a) (Tensor b -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor b
b)
    x <- broadcastTo dims a
    y <- broadcastTo dims b
    tensor dims (zipWith f (tensorValues x) (tensorValues y))

{- | Given two tensors, broadcast compatible shapes and add corresponding values, returning
a new contiguous tensor. Incompatible/overflowing broadcast shapes return Left
MathError.
-}
tensorAdd :: (Num a) => Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorAdd :: forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorAdd = (a -> a -> a)
-> Tensor a -> Tensor a -> Either MathError (Tensor a)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor a -> a -> a
forall a. Num a => a -> a -> a
(+)

{- | Given two tensors, broadcast compatible shapes and subtract right values from left,
returning a new contiguous tensor. Incompatible/overflowing broadcast shapes return
Left MathError.
-}
tensorSubtract :: (Num a) => Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorSubtract :: forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorSubtract = (a -> a -> a)
-> Tensor a -> Tensor a -> Either MathError (Tensor a)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor (-)

{- | Given two tensors, broadcast compatible shapes and multiply corresponding values,
returning a new contiguous tensor. Incompatible/overflowing broadcast shapes return
Left MathError.
-}
tensorMultiply :: (Num a) => Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorMultiply :: forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorMultiply = (a -> a -> a)
-> Tensor a -> Tensor a -> Either MathError (Tensor a)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor a -> a -> a
forall a. Num a => a -> a -> a
(*)

{- | Given two tensors, broadcast compatible shapes and divide left values by right,
returning a new contiguous tensor. Incompatible/overflowing broadcast shapes return
Left MathError. Division-by-zero/non-finite behavior follows Fractional; it is not
additionally checked.
-}
tensorDivide :: (Fractional a) => Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorDivide :: forall a.
Fractional a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorDivide = (a -> a -> a)
-> Tensor a -> Tensor a -> Either MathError (Tensor a)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor a -> a -> a
forall a. Fractional a => a -> a -> a
(/)

{- | Given two tensors, broadcast compatible shapes and compare corresponding values for
equality, returning a new contiguous tensor. Incompatible/overflowing broadcast shapes
return Left MathError.
-}
tensorEqual :: (Eq a) => Tensor a -> Tensor a -> Either MathError (Tensor Bool)
tensorEqual :: forall a.
Eq a =>
Tensor a -> Tensor a -> Either MathError (Tensor Bool)
tensorEqual = (a -> a -> Bool)
-> Tensor a -> Tensor a -> Either MathError (Tensor Bool)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)

{- | Given two tensors, broadcast compatible shapes and compare corresponding values with
(\<), returning a new contiguous tensor. Incompatible/overflowing broadcast shapes
return Left MathError.
-}
tensorLess :: (Ord a) => Tensor a -> Tensor a -> Either MathError (Tensor Bool)
tensorLess :: forall a.
Ord a =>
Tensor a -> Tensor a -> Either MathError (Tensor Bool)
tensorLess = (a -> a -> Bool)
-> Tensor a -> Tensor a -> Either MathError (Tensor Bool)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)

{- | Given two tensors, broadcast compatible shapes and combine Bool values with logical
AND, returning a new contiguous tensor. Incompatible/overflowing broadcast shapes
return Left MathError.
-}
tensorAnd :: Tensor Bool -> Tensor Bool -> Either MathError (Tensor Bool)
tensorAnd :: Tensor Bool -> Tensor Bool -> Either MathError (Tensor Bool)
tensorAnd = (Bool -> Bool -> Bool)
-> Tensor Bool -> Tensor Bool -> Either MathError (Tensor Bool)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor Bool -> Bool -> Bool
(&&)

{- | Given two tensors, broadcast compatible shapes and combine Bool values with logical OR,
returning a new contiguous tensor. Incompatible/overflowing broadcast shapes return
Left MathError.
-}
tensorOr :: Tensor Bool -> Tensor Bool -> Either MathError (Tensor Bool)
tensorOr :: Tensor Bool -> Tensor Bool -> Either MathError (Tensor Bool)
tensorOr = (Bool -> Bool -> Bool)
-> Tensor Bool -> Tensor Bool -> Either MathError (Tensor Bool)
forall a b c.
(a -> b -> c)
-> Tensor a -> Tensor b -> Either MathError (Tensor c)
zipTensor Bool -> Bool -> Bool
(||)

{- | Negate each Bool element, returning a new contiguous tensor of the same shape. No
broadcast or ownership mutation occurs.
-}
tensorNot :: Tensor Bool -> Tensor Bool
tensorNot :: Tensor Bool -> Tensor Bool
tensorNot = (Bool -> Bool) -> Tensor Bool -> Tensor Bool
forall a b. (a -> b) -> Tensor a -> Tensor b
mapTensor Bool -> Bool
not

{- | Explicitly convert each Real element through exact Rational to the destination
Fractional type, preserving shape. Destination rounding/overflow follows its
conversion; no implicit promotion occurs elsewhere.
-}
promoteTensor :: (Real a, Fractional b) => Tensor a -> Tensor b
promoteTensor :: forall a b. (Real a, Fractional b) => Tensor a -> Tensor b
promoteTensor = (a -> b) -> Tensor a -> Tensor b
forall a b. (a -> b) -> Tensor a -> Tensor b
mapTensor a -> b
forall a b. (Real a, Fractional b) => a -> b
realToFrac

reduceAxes :: (b -> a -> b) -> b -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor b)
reduceAxes :: forall b a.
(b -> a -> b)
-> b -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor b)
reduceAxes b -> a -> b
f b
initial [Int]
axes Bool
keep Tensor a
t
    | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub [Int]
axes) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
axes Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims) [Int]
axes =
        MathError -> Either MathError (Tensor b)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"reduction: invalid/duplicate axis")
    | Bool
otherwise = do
        let output :: [Int]
output =
                if Bool
keep
                    then [if Int
i Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
axes then Int
1 else Int
d | (Int
i, Int
d) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Int]
dims]
                    else [Int
d | (Int
i, Int
d) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Int]
dims, Int
i Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
axes]
            key :: [b] -> [b]
key [b]
ix =
                if Bool
keep
                    then [if Int
i Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
axes then b
0 else b
x | (Int
i, b
x) <- [Int] -> [b] -> [(Int, b)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [b]
ix]
                    else [b
x | (Int
i, b
x) <- [Int] -> [b] -> [(Int, b)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [b]
ix, Int
i Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
axes]
            accumulated :: Map [Int] b
accumulated =
                (Map [Int] b -> [Int] -> Map [Int] b)
-> Map [Int] b -> [[Int]] -> Map [Int] b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                    (\Map [Int] b
m [Int]
ix -> (Maybe b -> Maybe b) -> [Int] -> Map [Int] b -> Map [Int] b
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (b -> Maybe b
forall a. a -> Maybe a
Just (b -> Maybe b) -> (Maybe b -> b) -> Maybe b -> Maybe b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\b
old -> b -> a -> b
f b
old (Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex Tensor a
t [Int]
ix)) (b -> b) -> (Maybe b -> b) -> Maybe b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> (b -> b) -> Maybe b -> b
forall b a. b -> (a -> b) -> Maybe a -> b
maybe b
initial b -> b
forall a. a -> a
id) ([Int] -> [Int]
forall {a}. Num a => [a] -> [a]
key [Int]
ix) Map [Int] b
m)
                    Map [Int] b
forall k a. Map k a
M.empty
                    ([Int] -> [[Int]]
allIndices [Int]
dims)
        [Int] -> [b] -> Either MathError (Tensor b)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int]
output [b -> [Int] -> Map [Int] b -> b
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault b
initial [Int]
ix Map [Int] b
accumulated | [Int]
ix <- [Int] -> [[Int]]
allIndices [Int]
output]
  where
    dims :: [Int]
dims = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t

{- | Given unique zero-based axis list,keep-dimensions Bool,then tensor, reduce those axes
by sum. True retains reduced dimensions as size one; False removes them.
Invalid/duplicate axes fail. Empty reduced groups use zero.
-}
sumAxes :: (Num a) => [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
sumAxes :: forall a.
Num a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
sumAxes = (a -> a -> a)
-> a -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
forall b a.
(b -> a -> b)
-> b -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor b)
reduceAxes a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0

{- | Given unique zero-based axis list,keep-dimensions Bool,then tensor, reduce those axes
by product. True retains reduced dimensions as size one; False removes them.
Invalid/duplicate axes fail. Empty reduced groups use one.
-}
productAxes :: (Num a) => [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
productAxes :: forall a.
Num a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
productAxes = (a -> a -> a)
-> a -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
forall b a.
(b -> a -> b)
-> b -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor b)
reduceAxes a -> a -> a
forall a. Num a => a -> a -> a
(*) a
1

{- | Given unique zero-based axis list,keep-dimensions Bool,then tensor, reduce those axes
by arithmetic mean. True retains reduced dimensions as size one; False removes them.
Invalid/duplicate axes fail. Nonempty output with an empty reduced group fails.
-}
meanAxes :: (Fractional a) => [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
meanAxes :: forall a.
Fractional a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
meanAxes [Int]
axes Bool
keep Tensor a
t = do
    summed <- [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
forall a.
Num a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
sumAxes [Int]
axes Bool
keep Tensor a
t
    let count = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
i) | Int
i <- [Int]
axes]
    if count == 0 && numel summed > 0
        then Left (DomainError "meanAxes: empty reduction")
        else Right (mapTensor (/ fromInteger count) summed)

reduceNonempty :: (a -> a -> a) -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
reduceNonempty :: forall a.
(a -> a -> a)
-> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
reduceNonempty a -> a -> a
f [Int]
axes Bool
keep Tensor a
t = do
    reduced <- (Maybe a -> a -> Maybe a)
-> Maybe a
-> [Int]
-> Bool
-> Tensor a
-> Either MathError (Tensor (Maybe a))
forall b a.
(b -> a -> b)
-> b -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor b)
reduceAxes (\Maybe a
old a
x -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
x (a -> a -> a
`f` a
x) Maybe a
old)) Maybe a
forall a. Maybe a
Nothing [Int]
axes Bool
keep Tensor a
t
    case sequence (tensorValues reduced) of
        Maybe [a]
Nothing -> MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"reduction: empty axis")
        Just [a]
xs -> [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor (Tensor (Maybe a) -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor (Maybe a)
reduced) [a]
xs

{- | Given unique zero-based axis list,keep-dimensions Bool,then tensor, reduce those axes
by minimum. True retains reduced dimensions as size one; False removes them.
Invalid/duplicate axes fail. Empty reduced groups fail.
-}
minAxes :: (Ord a) => [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
minAxes :: forall a.
Ord a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
minAxes = (a -> a -> a)
-> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
forall a.
(a -> a -> a)
-> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
reduceNonempty a -> a -> a
forall a. Ord a => a -> a -> a
min

{- | Given unique zero-based axis list,keep-dimensions Bool,then tensor, reduce those axes
by maximum. True retains reduced dimensions as size one; False removes them.
Invalid/duplicate axes fail. Empty reduced groups fail.
-}
maxAxes :: (Ord a) => [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
maxAxes :: forall a.
Ord a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
maxAxes = (a -> a -> a)
-> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
forall a.
(a -> a -> a)
-> [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
reduceNonempty a -> a -> a
forall a. Ord a => a -> a -> a
max

argExtreme :: (Ord a) => (a -> a -> Bool) -> Tensor a -> Either MathError Int
argExtreme :: forall a.
Ord a =>
(a -> a -> Bool) -> Tensor a -> Either MathError Int
argExtreme a -> a -> Bool
better Tensor a
t = case [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t) of
    [] -> MathError -> Either MathError Int
forall a b. a -> Either a b
Left (String -> MathError
DomainError String
"argmin/argmax: empty tensor")
    (Int, a)
x : [(Int, a)]
xs -> Int -> Either MathError Int
forall a b. b -> Either a b
Right ((Int, a) -> Int
forall a b. (a, b) -> a
fst (((Int, a) -> (Int, a) -> (Int, a))
-> (Int, a) -> [(Int, a)] -> (Int, a)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\(Int, a)
old (Int, a)
new -> if a -> a -> Bool
better ((Int, a) -> a
forall a b. (a, b) -> b
snd (Int, a)
new) ((Int, a) -> a
forall a b. (a, b) -> b
snd (Int, a)
old) then (Int, a)
new else (Int, a)
old) (Int, a)
x [(Int, a)]
xs))

{- | Return first zero-based flat row-major index of a minimum logical element. Empty input
returns Left DomainError.
-}
argmin :: (Ord a) => Tensor a -> Either MathError Int
argmin :: forall a. Ord a => Tensor a -> Either MathError Int
argmin = (a -> a -> Bool) -> Tensor a -> Either MathError Int
forall a.
Ord a =>
(a -> a -> Bool) -> Tensor a -> Either MathError Int
argExtreme a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)

{- | Return first zero-based flat row-major index of a maximum logical element. Empty input
returns Left DomainError.
-}
argmax :: (Ord a) => Tensor a -> Either MathError Int
argmax :: forall a. Ord a => Tensor a -> Either MathError Int
argmax = (a -> a -> Bool) -> Tensor a -> Either MathError Int
forall a.
Ord a =>
(a -> a -> Bool) -> Tensor a -> Either MathError Int
argExtreme a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>)

{- | Given unique zero-based axis list,keep-dimensions Bool,then tensor, reduce those axes
by population variance (dividing by group size). True retains reduced dimensions as
size one; False removes them. Invalid/duplicate axes fail. Empty reduced groups fail as
meanAxes.
-}
varianceAxes :: (Fractional a) => [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
varianceAxes :: forall a.
Fractional a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
varianceAxes [Int]
axes Bool
keep Tensor a
t = do
    means <- [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
forall a.
Fractional a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
meanAxes [Int]
axes Bool
True Tensor a
t
    deviations <- zipTensor (\a
x a
m -> (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
m) a -> a -> a
forall a. Num a => a -> a -> a
* (a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
m)) t means
    meanAxes axes keep deviations

{- | Given unique zero-based axis list,keep-dimensions Bool,then tensor, reduce those axes
by scaled Euclidean norm. True retains reduced dimensions as size one; False removes
them. Invalid/duplicate axes fail. Empty groups have norm zero; IEEE NaN/infinity
propagate.
-}
normAxes :: (RealFloat a) => [Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
normAxes :: forall a.
RealFloat a =>
[Int] -> Bool -> Tensor a -> Either MathError (Tensor a)
normAxes [Int]
axes Bool
keep Tensor a
t = do
    pairs <- ((a, a) -> a -> (a, a))
-> (a, a)
-> [Int]
-> Bool
-> Tensor a
-> Either MathError (Tensor (a, a))
forall b a.
(b -> a -> b)
-> b -> [Int] -> Bool -> Tensor a -> Either MathError (Tensor b)
reduceAxes (a, a) -> a -> (a, a)
forall {b}. RealFloat b => (b, b) -> b -> (b, b)
accumulate (a
0, a
1) [Int]
axes Bool
keep Tensor a
t
    Right (mapTensor (\(a
scale, a
ssq) -> if a
scale a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then a
0 else a
scale a -> a -> a
forall a. Num a => a -> a -> a
* a -> a
forall a. Floating a => a -> a
sqrt a
ssq) pairs)
  where
    accumulate :: (b, b) -> b -> (b, b)
accumulate (b
scale, b
ssq) b
x
        | b -> Bool
forall a. RealFloat a => a -> Bool
isNaN b
x Bool -> Bool -> Bool
|| b -> Bool
forall a. RealFloat a => a -> Bool
isNaN b
scale = (b
0 b -> b -> b
forall a. Fractional a => a -> a -> a
/ b
0, b
0 b -> b -> b
forall a. Fractional a => a -> a -> a
/ b
0)
        | b -> Bool
forall a. RealFloat a => a -> Bool
isInfinite b
x Bool -> Bool -> Bool
|| b -> Bool
forall a. RealFloat a => a -> Bool
isInfinite b
scale = (b
1 b -> b -> b
forall a. Fractional a => a -> a -> a
/ b
0, b
1)
        | b
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
0 = (b
scale, b
ssq)
        | b
scale b -> b -> Bool
forall a. Ord a => a -> a -> Bool
< b -> b
forall a. Num a => a -> a
abs b
x = (b -> b
forall a. Num a => a -> a
abs b
x, b
1 b -> b -> b
forall a. Num a => a -> a -> a
+ b
ssq b -> b -> b
forall a. Num a => a -> a -> a
* (b
scale b -> b -> b
forall a. Fractional a => a -> a -> a
/ b -> b
forall a. Num a => a -> a
abs b
x) b -> Int -> b
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int))
        | Bool
otherwise = (b
scale, b
ssq b -> b -> b
forall a. Num a => a -> a -> a
+ (b -> b
forall a. Num a => a -> a
abs b
x b -> b -> b
forall a. Fractional a => a -> a -> a
/ b
scale) b -> Int -> b
forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
2 :: Int))

cumulative :: (a -> a -> a) -> a -> Int -> Tensor a -> Either MathError (Tensor a)
cumulative :: forall a.
(a -> a -> a)
-> a -> Int -> Tensor a -> Either MathError (Tensor a)
cumulative a -> a -> a
f a
initial Int
axis Tensor a
t = do
    Int -> Tensor a -> Either MathError ()
forall a. Int -> Tensor a -> Either MathError ()
axisValid Int
axis Tensor a
t
    [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor
        (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)
        [ (a -> a -> a) -> a -> [a] -> a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl a -> a -> a
f a
initial [Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex Tensor a
t (Int -> Int -> [Int] -> [Int]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
axis Int
k [Int]
ix) | Int
k <- [Int
0 .. [Int]
ix [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis]]
        | [Int]
ix <- [Int] -> [[Int]]
allIndices (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)
        ]

{- | Given existing axis then tensor, return inclusive running sums along that axis,
preserving shape. Invalid axis fails; arithmetic follows Num.
-}
cumulativeSum :: (Num a) => Int -> Tensor a -> Either MathError (Tensor a)
cumulativeSum :: forall a. Num a => Int -> Tensor a -> Either MathError (Tensor a)
cumulativeSum = (a -> a -> a)
-> a -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
(a -> a -> a)
-> a -> Int -> Tensor a -> Either MathError (Tensor a)
cumulative a -> a -> a
forall a. Num a => a -> a -> a
(+) a
0

{- | Given existing axis then tensor, return inclusive running products along that axis,
preserving shape. Invalid axis fails; arithmetic follows Num.
-}
cumulativeProduct :: (Num a) => Int -> Tensor a -> Either MathError (Tensor a)
cumulativeProduct :: forall a. Num a => Int -> Tensor a -> Either MathError (Tensor a)
cumulativeProduct = (a -> a -> a)
-> a -> Int -> Tensor a -> Either MathError (Tensor a)
forall a.
(a -> a -> a)
-> a -> Int -> Tensor a -> Either MathError (Tensor a)
cumulative a -> a -> a
forall a. Num a => a -> a -> a
(*) a
1

{- | Given contraction axes for the first tensor,axes for the second,then both tensors, sum
products over matched axis pairs. Noncontracted left axes precede right axes in the
output. Reject duplicate/out-of-range axes or mismatched contraction dimensions.
-}
tensordot :: (Num a) => [Int] -> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
tensordot :: forall a.
Num a =>
[Int]
-> [Int] -> Tensor a -> Tensor a -> Either MathError (Tensor a)
tensordot [Int]
axesA [Int]
axesB Tensor a
a Tensor a
b
    | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
axesA Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
axesB
        Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub [Int]
axesA) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
axesA
        Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub [Int]
axesB) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
axesB
        Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
da) [Int]
axesA
        Bool -> Bool -> Bool
|| (Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\Int
i -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
db) [Int]
axesB =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tensordot: invalid axes")
    | (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Int]
da [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
axesA [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
/= (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Int]
db [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
axesB =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tensordot: contracted dimensions differ")
    | Bool
otherwise = do
        _ <- [Int] -> Either MathError Int
shapeSize [Int]
output
        let na = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
freeA
            position a
axis [a]
contracted [a]
frees [a]
ix [a]
contraction = case a -> [a] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex a
axis [a]
contracted of
                Just Int
p -> [a]
contraction [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
p
                Maybe Int
Nothing -> case a -> [a] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex a
axis [a]
frees of Just Int
p -> [a]
ix [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
p; Maybe Int
Nothing -> String -> a
forall a. HasCallStack => String -> a
error String
"tensordot axis invariant"
        tensor
            output
            [ sum
                [ unsafeIndex a [position axis axesA freeA (take na ix) contracted | axis <- [0 .. length da - 1]]
                    * unsafeIndex b [position axis axesB freeB (drop na ix) contracted | axis <- [0 .. length db - 1]]
                | contracted <- allIndices (map (da !!) axesA)
                ]
            | ix <- allIndices output
            ]
  where
    da :: [Int]
da = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
a
    db :: [Int]
db = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
b
    freeA :: [Int]
freeA = [Int
i | Int
i <- [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
da Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
i Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
axesA]
    freeB :: [Int]
freeB = [Int
i | Int
i <- [Int
0 .. [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
db Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], Int
i Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
axesB]
    output :: [Int]
output = (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Int]
da [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
freeA [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map ([Int]
db [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!!) [Int]
freeB

{- | Explicit Einstein notation; each operand has one label per axis. Repeated
labels select diagonals; absent output labels are summed. Ellipses are not
accepted: callers express batch labels explicitly, preventing ambiguity.
-}
einsum :: (Num a) => [String] -> String -> [Tensor a] -> Either MathError (Tensor a)
einsum :: forall a.
Num a =>
[String] -> String -> [Tensor a] -> Either MathError (Tensor a)
einsum [String]
labels String
output [Tensor a]
operands
    | [String] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
labels Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Tensor a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Tensor a]
operands
        Bool -> Bool -> Bool
|| [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((String -> Tensor a -> Bool) -> [String] -> [Tensor a] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\String
ls Tensor a
t -> String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
ls Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)) [String]
labels [Tensor a]
operands)
        Bool -> Bool -> Bool
|| String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (ShowS
forall a. Eq a => [a] -> [a]
nub String
output) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
output =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"einsum: invalid label ranks or duplicate output labels")
    | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Char -> String -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` String
allLabels) String
output = MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"einsum: unknown output label")
    | ([(Char, Int)] -> Bool) -> [[(Char, Int)]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\[(Char, Int)]
ls -> ((Int, Int) -> Bool) -> [(Int, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(Int
x, Int
y) -> Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
y) [(Int
d, Char -> Int
dimension Char
l) | (Char
l, Int
d) <- [(Char, Int)]
ls]) [[(Char, Int)]]
labeledDims =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"einsum: label dimensions differ")
    | Bool
otherwise = do
        let dims :: [Int]
dims = (Char -> Int) -> String -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Int
dimension String
output; contracted :: String
contracted = (Char -> Bool) -> ShowS
forall a. (a -> Bool) -> [a] -> [a]
filter (Char -> String -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` String
output) String
allLabels
        _ <- [Int] -> Either MathError Int
shapeSize [Int]
dims
        tensor
            dims
            [ sum
                [ let environment = [(Char, Int)] -> Map Char Int
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList (String -> [Int] -> [(Char, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip String
output [Int]
ix [(Char, Int)] -> [(Char, Int)] -> [(Char, Int)]
forall a. [a] -> [a] -> [a]
++ String -> [Int] -> [(Char, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip String
contracted [Int]
ks)
                   in product [unsafeIndex t [environment M.! l | l <- ls] | (ls, t) <- zip labels operands]
                | ks <- allIndices (map dimension contracted)
                ]
            | ix <- allIndices dims
            ]
  where
    labeledDims :: [[(Char, Int)]]
labeledDims = (String -> Tensor a -> [(Char, Int)])
-> [String] -> [Tensor a] -> [[(Char, Int)]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\String
ls Tensor a
t -> String -> [Int] -> [(Char, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip String
ls (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t)) [String]
labels [Tensor a]
operands
    allLabels :: String
allLabels = ShowS
forall a. Eq a => [a] -> [a]
nub ([String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String]
labels)
    dimension :: Char -> Int
dimension Char
l = case [Int
d | [(Char, Int)]
ls <- [[(Char, Int)]]
labeledDims, (Char
name, Int
d) <- [(Char, Int)]
ls, Char
name Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
l] of Int
d : [Int]
_ -> Int
d; [] -> Int
0

{- | Multiply tensors on their final two dimensions and broadcast leading batch dimensions.
Require rank>=2 and matching inner dimensions; incompatible broadcast or size overflow
fails. Generic Num arithmetic is unchanged.
-}
batchedMatmul :: (Num a) => Tensor a -> Tensor a -> Either MathError (Tensor a)
batchedMatmul :: forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
batchedMatmul Tensor a
a Tensor a
b
    | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
da Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 Bool -> Bool -> Bool
|| [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
db Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"batchedMatmul: ranks must be at least two")
    | [Int]
da [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
da Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int]
db [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
db Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (Int -> Int -> MathError
InnerDimensionMismatch ([Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
da) ([Int]
db [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
db Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2)))
    | Bool
otherwise = do
        batch <- [Int] -> [Int] -> Either MathError [Int]
inferBroadcast (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
da Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) [Int]
da) (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
db Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) [Int]
db)
        let m = [Int]
da [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
da Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2); n = [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
db; k = [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
da
        av <- broadcastTo (batch ++ [m, k]) a
        bv <- broadcastTo (batch ++ [k, n]) b
        let output = [Int]
batch [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
m, Int
n]; rank = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
batch
        _ <- shapeSize output
        tensor
            output
            [ sum
                [ unsafeIndex av (take rank ix ++ [ix !! rank, l])
                    * unsafeIndex bv (take rank ix ++ [l, ix !! (rank + 1)])
                | l <- [0 .. k - 1]
                ]
            | ix <- allIndices output
            ]
  where
    da :: [Int]
da = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
a; db :: [Int]
db = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
b

{- | Return multidimensional Kronecker product, left-padding the lower-rank shape with ones
and multiplying corresponding axis sizes. Invalid/overflowing resulting dimensions
fail.
-}
tensorKronecker :: (Num a) => Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorKronecker :: forall a.
Num a =>
Tensor a -> Tensor a -> Either MathError (Tensor a)
tensorKronecker Tensor a
a Tensor a
b
    | (Integer -> Bool) -> [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)) [Integer]
dimensions =
        MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"tensorKronecker: dimension overflow")
    | Bool
otherwise = do
        av <- [Int] -> Tensor a -> Either MathError (Tensor a)
forall a. [Int] -> Tensor a -> Either MathError (Tensor a)
reshapeTensor [Int]
da Tensor a
a
        bv <- reshapeTensor db b
        let output = (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger [Integer]
dimensions
        _ <- shapeSize output
        tensor
            output
            [unsafeIndex av (zipWith div ix db) * unsafeIndex bv (zipWith mod ix db) | ix <- allIndices output]
  where
    rank :: Int
rank = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
a)) ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
b))
    pad :: [a] -> [a]
pad [a]
xs = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate (Int
rank Int -> Int -> Int
forall a. Num a => a -> a -> a
- [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) a
1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
xs
    da :: [Int]
da = [Int] -> [Int]
forall {a}. Num a => [a] -> [a]
pad (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
a)
    db :: [Int]
db = [Int] -> [Int]
forall {a}. Num a => [a] -> [a]
pad (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
b)
    dimensions :: [Integer]
dimensions = (Int -> Int -> Integer) -> [Int] -> [Int] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
x Int
y -> Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
x Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
y) [Int]
da [Int]
db

{- | Serialize (shape,logical row-major values) using Haskell Show. Storage strides/offsets
are not retained; the encoding requires a compatible Read instance for round-trip.
-}
serializeTensor :: (Show a) => Tensor a -> String
serializeTensor :: forall a. Show a => Tensor a -> String
serializeTensor Tensor a
t = ([Int], [a]) -> String
forall a. Show a => a -> String
show (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t, Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t)

{- | Read a complete Haskell (shape,values) encoding, permitting only trailing whitespace,
and validate shape/count. Malformed input or invalid dimensions/count returns Left
MathError.
-}
deserializeTensor :: (Read a) => String -> Either MathError (Tensor a)
deserializeTensor :: forall a. Read a => String -> Either MathError (Tensor a)
deserializeTensor String
input = case ReadS ([Int], [a])
forall a. Read a => ReadS a
reads String
input of
    [(([Int]
dims, [a]
xs), String
rest)] | (Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Char -> String -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` String
" \n\r\t") String
rest -> [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int]
dims [a]
xs
    [(([Int], [a]), String)]
_ -> MathError -> Either MathError (Tensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"deserializeTensor: invalid shape/value encoding")

{- | Lazy expression trees defer computation. Evaluation fuses adjacent maps
and composes per-element expressions through broadcasted binary nodes; only
the output array is allocated. Shape errors are checked before evaluation.
-}
data TensorGraph a = Graph [Int] ([Int] -> a)

{- | Wrap an immutable tensor as a lazy TensorGraph leaf, retaining its backing storage
without evaluating a new output array.
-}
inputGraph :: Tensor a -> TensorGraph a
inputGraph :: forall a. Tensor a -> TensorGraph a
inputGraph Tensor a
t = [Int] -> ([Int] -> a) -> TensorGraph a
forall a. [Int] -> ([Int] -> a) -> TensorGraph a
Graph (Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t) (Tensor a -> [Int] -> a
forall a. Tensor a -> [Int] -> a
unsafeIndex Tensor a
t)

{- | Given element function then graph, compose a lazy element transform, preserving shape.
No output tensor is allocated until evaluateGraph.
-}
mapGraph :: (a -> b) -> TensorGraph a -> TensorGraph b
mapGraph :: forall a b. (a -> b) -> TensorGraph a -> TensorGraph b
mapGraph a -> b
f (Graph [Int]
dims [Int] -> a
action) = [Int] -> ([Int] -> b) -> TensorGraph b
forall a. [Int] -> ([Int] -> a) -> TensorGraph a
Graph [Int]
dims (a -> b
f (a -> b) -> ([Int] -> a) -> [Int] -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> a
action)

{- | Given binary function then two graphs, validate broadcast shapes and return a fused
lazy graph. Incompatible shapes return Left MathError before evaluation.
-}
zipGraph :: (a -> b -> c) -> TensorGraph a -> TensorGraph b -> Either MathError (TensorGraph c)
zipGraph :: forall a b c.
(a -> b -> c)
-> TensorGraph a
-> TensorGraph b
-> Either MathError (TensorGraph c)
zipGraph a -> b -> c
f (Graph [Int]
da [Int] -> a
a) (Graph [Int]
db [Int] -> b
b) = do
    dims <- [Int] -> [Int] -> Either MathError [Int]
inferBroadcast [Int]
da [Int]
db
    let select [a]
source [c]
ix = (a -> c -> c) -> [a] -> [c] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
d c
i -> if a
d a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 then c
0 else c
i) [a]
source (Int -> [c] -> [c]
forall a. Int -> [a] -> [a]
drop ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Int
forall a. Num a => a -> a -> a
- [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
source) [c]
ix)
    Right (Graph dims (\[Int]
ix -> a -> b -> c
f ([Int] -> a
a ([Int] -> [Int] -> [Int]
forall {a} {c}. (Eq a, Num a, Num c) => [a] -> [c] -> [c]
select [Int]
da [Int]
ix)) ([Int] -> b
b ([Int] -> [Int] -> [Int]
forall {a} {c}. (Eq a, Num a, Num c) => [a] -> [c] -> [c]
select [Int]
db [Int]
ix))))

{- | Evaluate a lazy graph into one contiguous row-major tensor. Adjacent map/zip operations
are fused; scalar evaluation behavior follows supplied functions and checked output
shape errors propagate.
-}
evaluateGraph :: TensorGraph a -> Either MathError (Tensor a)
evaluateGraph :: forall a. TensorGraph a -> Either MathError (Tensor a)
evaluateGraph (Graph [Int]
dims [Int] -> a
action) = [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
tensor [Int]
dims [[Int] -> a
action [Int]
ix | [Int]
ix <- [Int] -> [[Int]]
allIndices [Int]
dims]