0 | {--
   1 | Copyright (C) 2021  Joel Berkeley
   2 |
   3 | This program is free software: you can redistribute it and/or modify
   4 | it under the terms of the GNU Affero General Public License as published
   5 | by the Free Software Foundation, either version 3 of the License, or
   6 | (at your option) any later version.
   7 |
   8 | This program is distributed in the hope that it will be useful,
   9 | but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 | GNU Affero General Public License for more details.
  12 |
  13 | You should have received a copy of the GNU Affero General Public License
  14 | along with this program.  If not, see <https://www.gnu.org/licenses/>.
  15 | --}
  16 | ||| Defines `Tensor`, an array of numbers or booleans, along with a number of functions operating on
  17 | ||| `Tensor`s. `Tensor` operations typically leverage hardware acceleration and graph compilation.
  18 | ||| spidr tracks tensor shape and data type in the types, so you can be sure that if your tensor
  19 | ||| code compiles, these are consistent.
  20 | |||
  21 | ||| spidr achieves efficient reuse of tensor computations with `Tag`. See the tutorial
  22 | ||| _Nuisances in the Tensor API_ for a discussion of pitfalls to avoid when using `Tag`.
  23 | module Tensor
  24 |
  25 | import public Control.Monad.State
  26 | import Control.Monad.Error.Either
  27 | import Syntax.PreorderReasoning
  28 |
  29 | import Compiler.Eval
  30 | import Compiler.IR
  31 | import Compiler.Array
  32 | import Compiler.LiteralRW
  33 | import Compiler.Passes
  34 | import Device
  35 | import public DType
  36 | import public Literal
  37 | import public Types
  38 | import public Util
  39 |
  40 | ||| A scalar or array. Construct a `Tensor` with function `tensor`.
  41 | export
  42 | data Tensor : Shape -> DType -> Type where
  43 |   MkTensor : Value -> {shape : _} -> {dtype : _} -> Tensor shape dtype
  44 |
  45 | (.type) : Tensor shape dtype -> ValueType
  46 | (.type) (MkTensor {shape, dtype} _) = TensorType shape dtype
  47 |
  48 | (.value) : Tensor s t -> Value
  49 | (.value) (MkTensor v) = v
  50 |
  51 | ||| The effect of tagging nodes in a computational graph.
  52 | export
  53 | data TagT : (Type -> Type) -> Type -> Type where
  54 |   MkTagT : StateT Env m a -> TagT m a
  55 |
  56 | public export 0
  57 | Tag : Type -> Type
  58 | Tag = TagT Identity
  59 |
  60 | export
  61 | Functor m => Functor (TagT m) where
  62 |   map f (MkTagT x) = MkTagT (map f x)
  63 |
  64 | export
  65 | Monad m => Applicative (TagT m) where
  66 |   pure x = MkTagT (pure x)
  67 |   (MkTagT f) <*> (MkTagT x) = MkTagT (f <*> x)
  68 |
  69 | export
  70 | Monad m => Monad (TagT m) where
  71 |   (MkTagT x) >>= f = MkTagT $ x >>= (\y => let MkTagT z = f y in z)
  72 |
  73 | export
  74 | MonadTrans TagT where
  75 |   lift = MkTagT . lift
  76 |
  77 | public export
  78 | interface Taggable a where
  79 |   ||| Mark an expression to be efficiently reused. For example, in
  80 |   ||| ```
  81 |   ||| bad : Tensor [9999999] F64
  82 |   ||| bad = let x = fill {shape = [9999999]} 1.0 in x + x
  83 |   |||
  84 |   ||| good : Tag $ Tensor [9999999] F64
  85 |   ||| good = do x <- tag $ fill {shape = [9999999]} 1.0
  86 |   |||           pure (x + x)
  87 |   ||| ```
  88 |   ||| the large vector `x` is calculated twice in `bad`, but once in `good`, as `tag` marks it for
  89 |   ||| sharing.
  90 |   |||
  91 |   ||| Types that implement this interface should `tag` constituent components it deems worth sharing.
  92 |   ||| For example, see the implementation for tuples.
  93 |   |||
  94 |   ||| See tutorial _Nuisances in the Tensor API_ for details.
  95 |   tag : Monad m => a -> TagT m a
  96 |
  97 | Taggable OpRef where
  98 |   tag = MkTagT . tagOpRef
  99 |
 100 | export
 101 | Taggable (Tensor shape dtype) where
 102 |   tag (MkTensor $ V idx op) = map (\op => MkTensor $ V idx op) (tag op)
 103 |
 104 | export
 105 | (Taggable a, Taggable b) => Taggable (a, b) where
 106 |   tag (a, b) = [| (tag a, tag b) |]
 107 |
 108 | val0 : Op -> Value
 109 | val0 = V 0 . Concrete
 110 |
 111 | t0 : {shape : _} -> {dtype : _} -> Op -> Tensor shape dtype
 112 | t0 x = MkTensor $ val0 x
 113 |
 114 | ||| Construct a `Tensor` from `Literal` data. For example
 115 | ||| ```
 116 | ||| x : Tensor [2, 3] S32
 117 | ||| x = tensor [[1, 2, 3],
 118 | |||             [4, 5, 6]]
 119 | ||| ```
 120 | export
 121 | tensor : {shape : _} -> {dtype : _} -> Literal shape (idrisType dtype) -> Tensor shape dtype
 122 | tensor lit = t0 $ Lit shape dtype lit
 123 |
 124 | namespace F64
 125 |   export
 126 |   fromDouble : Double -> Tensor [] F64
 127 |   fromDouble = tensor . Scalar
 128 |
 129 | try : Show e => EitherT e IO a -> IO a
 130 | try = eitherT (\e => assert_total $ idris_crash $ show e) pure
 131 |
 132 | %hide Literal.All2.All2
 133 |
 134 | namespace List
 135 |   namespace Tag
 136 |     ||| Evaluate a list of `Tensor`s as a list of `Literal`s. Tensors in the list can have different
 137 |     ||| shapes and element types. For example,
 138 |     ||| ```
 139 |     ||| main : Device -> IO ()
 140 |     ||| main device = do [x, y] <- eval device $ do let x = tensor {dtype = F64} [1.2, 3.4]
 141 |     |||                                             y <- reduce @{Sum} [0] x
 142 |     |||                                             pure [x, y]
 143 |     |||                  printLn x
 144 |     |||                  printLn y
 145 |     ||| ```
 146 |     ||| In contrast to `Tensor.eval` when called on multiple tensors, this function constructs and
 147 |     ||| compiles the graph just once.
 148 |     export covering
 149 |     eval :
 150 |       forall shapes, dtypes .
 151 |       Device ->
 152 |       Tag (All2 Tensor shapes dtypes) ->
 153 |       IO (All2 Literal shapes (DType.idrisType <$> dtypes))
 154 |     eval device (MkTagT xs) =
 155 |       let (env, xs) = runState empty xs
 156 |           main = MkFn 0 [] (extract (.type) xs) (extract (.value) xs) env
 157 |           types = mapProperty (\(MkTensor _) => MkTensorData ()) xs
 158 |        in try $ readAll <$> execute device main types
 159 |
 160 |       where
 161 |
 162 |       extract : (forall s, t . Tensor s t -> a) -> All2 Tensor ss tt -> Vect (length ss) a
 163 |       extract f [] = []
 164 |       extract f (x :: xs) = f x :: extract f xs
 165 |
 166 |       readAll :
 167 |         All2 (TensorData $ \_ => Array . DType.idrisType) s t ->
 168 |         All2 Literal s (DType.idrisType <$> t)
 169 |       readAll [] = []
 170 |       readAll (MkTensorData {dtype} l :: ls) = read dtype l :: readAll ls
 171 |
 172 |   ||| A convenience wrapper for `List.Tag.eval`, for use with a bare list of `Tensor`s.
 173 |   export covering
 174 |   eval :
 175 |     forall shapes, dtypes .
 176 |     Device ->
 177 |     All2 Tensor shapes dtypes ->
 178 |     IO (All2 Literal shapes (DType.idrisType <$> dtypes))
 179 |   eval device xs = eval device (pure xs)
 180 |
 181 | namespace Tag
 182 |   ||| Evaluate a `Tensor`, returning its value as a `Literal`. This function builds and executes the
 183 |   ||| computational graph.
 184 |   |||
 185 |   ||| **Note:** Each call to `eval` will rebuild and execute the graph; multiple calls to `eval` on
 186 |   ||| different tensors, even if they are in the same computation, will be treated independently.
 187 |   ||| To efficiently evaluate multiple tensors at once, use `List.Tag.eval`.
 188 |   export covering
 189 |   eval : Device -> Tag (Tensor shape dtype) -> IO (Literal shape (DType.idrisType dtype))
 190 |   eval device x = map (\[z] => z) $ List.Tag.eval device $ map (\z => [z]) x
 191 |
 192 | ||| A convenience wrapper for `Tag.eval`, for use with a bare `Tensor`.
 193 | export covering
 194 | eval : Device -> Tensor shape dtype -> IO (Literal shape (DType.idrisType dtype))
 195 | eval device x = eval device (pure x)
 196 |
 197 | ||| A string representation of a tensor graph.
 198 | |||
 199 | ||| There are no guarantees whatsoever as to the string structure and contents.
 200 | export
 201 | Show (Tag $ Tensor shape dtype) where
 202 |   show (MkTagT x) =
 203 |     let (env, MkTensor x) = runState empty x
 204 |      in show (MkFn 0 [] [TensorType shape dtype] [x] env)
 205 |
 206 | export
 207 | Show (Tensor shape dtype) where show = show . pure {f = Tag}
 208 |
 209 | ||| Positive infinity.
 210 | export
 211 | inf : Tensor [] F64
 212 |
 213 | ||| NaN (not a number).
 214 | export
 215 | nan : Tensor [] F64
 216 |
 217 | namespace Bound
 218 |   ||| Compares less than or equal to any other value (except NaN).
 219 |   export
 220 |   min : Ord dtype => Tensor [] dtype
 221 |   min @{OrdS32} = t0 $ MinValue S32
 222 |   min @{OrdS64} = t0 $ MinValue S64
 223 |   min @{OrdU32} = t0 $ MinValue U32
 224 |   min @{OrdU64} = t0 $ MinValue U64
 225 |   min @{OrdF64} = tensor $ Scalar $ -1.0 / 0.0
 226 |
 227 |   ||| Compares greater than or equal to any other value (except NaN).
 228 |   export
 229 |   max : Ord dtype => Tensor [] dtype
 230 |   max @{OrdS32} = t0 $ MaxValue S32
 231 |   max @{OrdS64} = t0 $ MaxValue S64
 232 |   max @{OrdU32} = t0 $ MaxValue U32
 233 |   max @{OrdU64} = t0 $ MaxValue U64
 234 |   max @{OrdF64} = tensor $ Scalar $ 1.0 / 0.0
 235 |
 236 | ||| The most negative possible finite float, approx. -1.8e308
 237 | export
 238 | minFinite : Tensor [] F64
 239 | minFinite = t0 MinFiniteFloat
 240 |
 241 | ||| The most positive possible finite float, approx. 1.8e308
 242 | export
 243 | maxFinite : Tensor [] F64
 244 | maxFinite = t0 MaxFiniteFloat
 245 |
 246 | ||| Cast the element type. For example, `castDtype (tensor {dtype = S32} [1, -2])` is
 247 | ||| `tensor {dtype = F64} [1.0, -2.0]`.
 248 | export
 249 | castDtype : Integral dtype => Tensor shape dtype -> Tensor shape F64
 250 | castDtype $ MkTensor {shape} x = t0 $ Convert F64 shape x
 251 |
 252 | ||| A function type. For example `Func [Nat, String] Bool` is `Nat -> String -> Bool`.
 253 | public export 0
 254 | Func : Vect arity Type -> Type -> Type
 255 | Func [] r = r
 256 | Func (t :: ts) r = t -> Func ts r
 257 |
 258 | mapFunc : (a -> b) -> {arity : Nat} -> {0 xs : Vect arity Type} -> Func xs a -> Func xs b
 259 | mapFunc f {arity = 0} {xs = []} g = f g
 260 | mapFunc f {arity = (S k)} {xs = _ :: _} g = \x => mapFunc f (g x)
 261 |
 262 | mkFn' :
 263 |   forall arity, rshapes, rdtypes . (shapes : Vect arity Shape) -> (dtypes : Vect arity DType) ->
 264 |   Func [| Tensor shapes dtypes |] (Tag $ All2 Tensor rshapes rdtypes) ->
 265 |   Tag (Fn arity, All2 Tensor rshapes rdtypes)
 266 | mkFn' shapes dtypes f = MkTagT $ do
 267 |   addr <- reserve
 268 |
 269 |   let leq : length shapes === arity
 270 |       leq = lengthCorrect shapes
 271 |
 272 |   let MkTagT res = applyNary addr (length shapes) _ _ $ rewrite leq in f
 273 |       (env, results) = runState (emptyFrom !get) res
 274 |       resAndTys = fromList $ resultsAndTypes results
 275 |       argValueTypes = zipWith TensorType shapes dtypes
 276 |       f : Fn arity = MkFn addr argValueTypes (snd <$> resAndTys) (fst <$> resAndTys) env
 277 |
 278 |   updateCounterFrom env
 279 |
 280 |   pure (f, results)
 281 |
 282 |   where
 283 |
 284 |   applyNary :
 285 |     (addr, ar : Nat) ->
 286 |     (ss : Vect ar Shape) ->
 287 |     (ds : Vect ar DType) ->
 288 |     Func [| Tensor ss ds |] a -> a
 289 |   applyNary addr 0 [] [] f = f
 290 |   applyNary addr (S a) (s :: ss) (t :: ts) f =
 291 |     applyNary addr a ss ts (f $ MkTensor $ V (length shapes `minus` S a) $ BoundSet addr)
 292 |
 293 |   resultsAndTypes : All2 Tensor s d -> List (Value, ValueType)
 294 |   resultsAndTypes [] = []
 295 |   resultsAndTypes ((MkTensor {shape, dtype} x) :: xs) =
 296 |     (x, TensorType shape dtype) :: resultsAndTypes xs
 297 |
 298 | mkFn :
 299 |   {arity : _} -> (shapes : Vect arity Shape) -> (dtypes : Vect arity DType) ->
 300 |   forall rshapes, rdtypes .
 301 |   Func [| Tensor shapes dtypes |] (Tag $ All2 Tensor rshapes rdtypes) -> Tag $ Fn arity
 302 | mkFn shapes dtypes f = fst <$> mkFn' shapes dtypes f
 303 |
 304 | mkFn1 :
 305 |   {arity : _} -> (shapes : Vect arity Shape) -> (dtypes : Vect arity DType) ->
 306 |   Func [| Tensor shapes dtypes |] (Tag $ Tensor rshape rdtype) -> Tag $ Fn arity
 307 | mkFn1 shapes dtypes f =
 308 |   mkFn shapes dtypes $ mapFunc (\x => pure $ the (All2 _ _ _) [!x]) f
 309 |
 310 | covering
 311 | namedFunc :
 312 |   {arity : _} -> {shapes : Vect arity Shape} -> {dtypes : Vect arity DType} ->
 313 |   forall rshapes, rdtypes .
 314 |   (shapeFn : Shape -> Shape) ->
 315 |   (forall extra . Fn (arity + extra) -> Vect arity Value -> Vect extra Value -> Op) ->
 316 |   Func [| Tensor shapes dtypes |] (Tag $ All2 Tensor rshapes rdtypes) ->
 317 |   Tag $ Func [| Tensor (shapeFn <$> shapes) dtypes |] $
 318 |     Tag $ All2 Tensor (shapeFn <$> rshapes) rdtypes
 319 | namedFunc shapeFn remapResults f = do
 320 |   (f, results) <- mkFn' shapes dtypes {rshapes, rdtypes} f
 321 |   let (extraArgs, f) = snd $ removeCaptures {n = arity} f
 322 |   MkTagT $ modify {ops $= ((f.tag, NamedFunc f) ::)}
 323 |   let remap = \args =>
 324 |         retarget 0 results <$> tag (Concrete $ remapResults f args extraArgs)
 325 |   pure $ mapFunc remap $ args [] arity (shapeFn <$> shapes) dtypes
 326 |
 327 |   where
 328 |
 329 |   args :
 330 |     (acc : Vect p Value) ->
 331 |     (0 ar : Nat) ->
 332 |     (ss : Vect ar Shape) ->
 333 |     (ds : Vect ar DType) ->
 334 |     Func [| Tensor ss ds |] $ Vect (p + ar) Value
 335 |   args acc 0 [] [] = rewrite plusZeroRightNeutral p in acc
 336 |   args acc (S a) (s :: ss) (t :: ts) = \(MkTensor x) =>
 337 |     rewrite sym $ plusSuccRightSucc p a in args (snoc acc x) a ss ts
 338 |
 339 |   retarget : Nat -> All2 Tensor ss ds -> OpRef -> All2 Tensor (shapeFn <$> ss) ds
 340 |   retarget k [] x = []
 341 |   retarget k (MkTensor {shape, dtype} _ :: xs) x =
 342 |     MkTensor {shape = shapeFn shape, dtype} (V k x) :: retarget (S k) xs x
 343 |
 344 | ||| Function abstraction in the framework IR.
 345 | |||
 346 | ||| Operations in spidr are, by default, inlined, which can lead to large IRs. `func` abstracts
 347 | ||| (and names) a function in the IR. The resulting function will have the exact same semantics as
 348 | ||| the input, though may perform differently, depending on ML compiler behaviour.
 349 | |||
 350 | ||| Implementation note: MLIR (more specifically the "func" dialect) only supports named functions
 351 | ||| at the top (or module) level. To achieve this, `func` lifts its function to the top level,
 352 | ||| and automatically converts any variable capture to new function arguments. As such, the IR
 353 | ||| may appear different to the Idris code from which it derives.
 354 | export covering
 355 | func :
 356 |   {shapes : Vect arity Shape} -> {dtypes : Vect arity DType} -> forall rshapes, rdtypes .
 357 |   let fn = Func [| Tensor shapes dtypes |] $ Tag $ All2 Tensor rshapes rdtypes in fn -> Tag fn
 358 | func f =
 359 |   rewrite sym $ functorIdentity shapes in
 360 |   rewrite sym $ functorIdentity rshapes in
 361 |   let lc = lengthCorrect (map id shapes) in
 362 |   rewrite sym lc in
 363 |     namedFunc {arity = length (map id shapes)} id (rewrite lc in remapResults) (rewrite lc in f)
 364 |
 365 |   where
 366 |   remapResults : Fn (arity + extra) -> Vect arity Value -> Vect extra Value -> Op
 367 |   remapResults f explicitArgs extraArgs =
 368 |     CallByName f.tag (toList $ f.resultTypes) (toList $ explicitArgs ++ extraArgs)
 369 |
 370 | ||| (Experimental) function vectorization.
 371 | |||
 372 | ||| Lift a function on tensors, so that it applies element-wise across the common leading dimension
 373 | ||| of its tensor arguments. For example, for
 374 | ||| ```
 375 | ||| xs : Tensor [2, 3, 3] S32
 376 | ||| xs = tensor [[[ 0,  1,  2],
 377 | |||               [ 3,  4,  5],
 378 | |||               [ 6,  7,  8]],
 379 | |||              [[ 9, 10, 11],
 380 | |||               [12, 13, 14],
 381 | |||               [15, 16, 17]]]
 382 | |||
 383 | ||| ys : Tensor [2] S32
 384 | ||| ys = Tensor [2, -1]
 385 | ||| ```
 386 | ||| `do !(vmap (\x, y => pure $ y * diag x)) xs ys` produces `tensor [[0, 8, 16], [-9, -13, -17]]`.
 387 | |||
 388 | ||| **Warning:** `vmap` is experimental, and only implemented for a subset of the tensor API. You
 389 | ||| can see approximately which operations are supported on the Enzyme [tracking issue](https://github.com/EnzymeAD/Enzyme-JAX/issues/152).
 390 | export partial
 391 | vmap :
 392 |   {n : _} ->
 393 |   {shapes : Vect arity Shape} -> {dtypes : Vect arity DType} ->
 394 |   {rshapes : _} -> {rdtypes : _} ->
 395 |   Func [| Tensor shapes dtypes |] (Tag $ All2 Tensor rshapes rdtypes) ->
 396 |   Tag $ Func [| Tensor (Prelude.map (n ::) shapes) dtypes |]
 397 |     (Tag $ All2 Tensor (Prelude.map (n ::) rshapes) rdtypes)
 398 | vmap f =
 399 |   let lc = lengthCorrect shapes in
 400 |   rewrite sym lc in
 401 |     namedFunc {arity = length shapes} (n ::) (rewrite lc in remapResults) (rewrite lc in f)
 402 |
 403 |   where
 404 |   remapResults : Fn (arity + extra) -> Vect arity Value -> Vect extra Value -> Op
 405 |   remapResults f explicitArgs extraArgs =
 406 |     let extraArgs = Prelude.map (V 0 . Concrete . Broadcast (AddLeading [n])) extraArgs
 407 |         outTys = zipWith (TensorType . (n ::)) rshapes rdtypes
 408 |      in Vectorize outTys [n] f.tag (toList $ explicitArgs ++ extraArgs)
 409 |
 410 | ||| (Experimental) reverse-mode automatic differentiation.
 411 | |||
 412 | ||| `grad` can be applied repeatedly to obtain higher derivatives, though we do not yet support
 413 | ||| derivatives of vector-valued functions.
 414 | |||
 415 | ||| For example, for
 416 | ||| ```
 417 | ||| f : Tensor [2] F64 -> Tag $ Tensor [] F64
 418 | ||| f x = do
 419 | |||   x <- tag x
 420 | |||   let (x0, x1) = (slice [at 0] x, slice [at 1] x)
 421 | |||   pure $ x0 / x1
 422 | ||| ```
 423 | ||| `grad f (tensor [3.0, 2.0])` produces `tensor [0.5, -0.75]`.
 424 | |||
 425 | ||| **Warning:** `grad` is experimental, and only implemented for a subset of the tensor API. You
 426 | ||| can see approximately which operations are supported on the Enzyme [tracking issue](https://github.com/EnzymeAD/Enzyme-JAX/issues/88).
 427 | export partial
 428 | grad : (Tensor shape F64 -> Tag $ Tensor [] F64) -> Tensor shape F64 -> Tag $ Tensor shape F64
 429 | grad f (MkTensor x) = pure $ t0 $ Grad shape !(mkFn1 [shape] [_] f) x
 430 |
 431 | ||| Reshape a `Tensor`. For example, `reshape {to = [2, 1]} (tensor [3, 4])` is
 432 | ||| `tensor [[3], [4]]`. The output can have a different rank to the input.
 433 | export
 434 | reshape :
 435 |   {to : _} ->
 436 |   {auto 0 sizesEqual : product from = product to} ->
 437 |   Tensor from dtype ->
 438 |   Tensor to dtype
 439 | reshape $ MkTensor {shape} x = t0 $ Reshape dtype to x
 440 |
 441 | ||| Add a dimension of length one at the specified `axis`. The new dimension will be at the
 442 | ||| specified `axis` in the new `Tensor` (as opposed to the original `Tensor`). For example,
 443 | ||| `expand 1 $ tensor [[1, 2], [3, 4], [5, 6]]` is `tensor [[[1, 2]], [[3, 4]], [[5, 6]]]`.
 444 | export
 445 | expand :
 446 |   (axis : Nat) ->
 447 |   {auto 0 inBounds : axis `LTE` length shape} ->
 448 |   Tensor shape dtype ->
 449 |   Tensor (insertAt axis 1 shape) dtype
 450 | expand axis $ MkTensor {shape = _} x = t0 $ Reshape dtype (insertAt axis 1 shape) x
 451 |
 452 | namespace Squeezable
 453 |   ||| A `Squeezable from to` constitutes proof that the shape `from` can be squeezed to the
 454 |   ||| shape `to`. Squeezing is the process of removing any number of dimensions of length one.
 455 |   public export
 456 |   data Squeezable : (0 from : Shape) -> (0 to : Shape) -> Type where
 457 |     ||| Proof that a shape can be squeezed to itself. For example:
 458 |     |||
 459 |     ||| [] to []
 460 |     ||| [3, 4] to [3, 4]
 461 |     Same : Squeezable x x
 462 |
 463 |     ||| Proof that any dimensions (including those of length 1) can be preserved in the process of
 464 |     ||| squeezing. For example:
 465 |     |||
 466 |     ||| ...
 467 |     Match : Squeezable from to -> Squeezable (x :: from) (x :: to)
 468 |
 469 |     ||| Proof that any dimensions of length one can be squeezed out. For example:
 470 |     |||
 471 |     ||| [1, 3, 1, 1, 4] to [3, 4]
 472 |     Nest : Squeezable from to -> Squeezable (1 :: from) to
 473 |
 474 | ||| Remove dimensions of length one from a `Tensor` such that it has the desired shape. For example:
 475 | |||
 476 | ||| ```
 477 | ||| x : Tensor [2, 1, 3, 1] S32
 478 | ||| x = tensor [[[[4], [5], [6]]],
 479 | |||             [[[7], [8], [9]]]]
 480 | |||
 481 | ||| y : Tensor [2, 1, 3] S32
 482 | ||| y = squeeze x
 483 | ||| ```
 484 | ||| is
 485 | ||| ```
 486 | ||| y : Tensor [2, 1, 3] S32
 487 | ||| y = tensor [[[4, 5, 6]],
 488 | |||             [[7, 8, 9]]]
 489 | ||| ```
 490 | export
 491 | squeeze :
 492 |   {to : _} ->
 493 |   {auto 0 shapesSqueezable : Squeezable from to} ->
 494 |   Tensor from dtype ->
 495 |   Tensor to dtype
 496 | squeeze $ MkTensor {shape} x = t0 $ Reshape dtype to x
 497 |
 498 | ||| A `SliceOrIndex d` is a valid slice or index into a dimension of size `d`. See `slice` for
 499 | ||| details.
 500 | export
 501 | data SliceOrIndex : Nat -> Type where
 502 |   Slice :
 503 |     (from, to : Nat) ->
 504 |     {size : _} ->
 505 |     {auto 0 fromTo : from + size = to} ->
 506 |     {auto 0 inDim : LTE to d} ->
 507 |     SliceOrIndex d
 508 |   Index : (idx : Nat) -> {auto 0 inDim : LT idx d} -> SliceOrIndex d
 509 |   DynamicSlice : Tensor [] U64 -> (size : Nat) -> {auto 0 inDim : LTE size d} -> SliceOrIndex d
 510 |   DynamicIndex : Tensor [] U64 -> SliceOrIndex d
 511 |
 512 | ||| Index at `idx`. See `slice` for details.
 513 | public export
 514 | at : (idx : Nat) -> {auto 0 inDim : LT idx d} -> SliceOrIndex d
 515 | at = Index
 516 |
 517 | namespace Dynamic
 518 |   ||| Index at the specified index. See `slice` for details.
 519 |   public export
 520 |   at : Tensor [] U64 -> SliceOrIndex d
 521 |   at = DynamicIndex
 522 |
 523 | ||| Slice from `from` (inclusive) to `to` (exclusive). See `slice` for details.
 524 | public export
 525 | (.to) :
 526 |   (from, to : Nat) ->
 527 |   {size : _} ->
 528 |   {auto 0 fromTo : from + size = to} ->
 529 |   {auto 0 inDim : LTE to d} ->
 530 |   SliceOrIndex d
 531 | (.to) = Slice
 532 |
 533 | ||| Slice `size` elements starting at the specified scalar `U64` index. See `slice` for details.
 534 | public export
 535 | (.size) : Tensor [] U64 -> (size : Nat) -> {auto 0 inDim : LTE size d} -> SliceOrIndex d
 536 | (.size) = DynamicSlice
 537 |
 538 | ||| Slice across all indices along an axis. See `slice` for details.
 539 | public export
 540 | all : {d : _} -> SliceOrIndex d
 541 | all = Slice 0 @{%search} @{reflexive {ty = Nat}} d
 542 |
 543 | ||| A `MultiSlice shape` is a valid multi-dimensional slice into a tensor with shape `shape`.
 544 | ||| See `slice` for details.
 545 | public export
 546 | data MultiSlice : Shape -> Type where
 547 |   Nil : MultiSlice ds
 548 |   (::) : SliceOrIndex d -> MultiSlice ds -> MultiSlice (d :: ds)
 549 |
 550 | namespace MultiSlice
 551 |   ||| The shape of a tensor produced by slicing with the specified multi-dimensional slice. See
 552 |   ||| `Tensor.slice` for details.
 553 |   public export
 554 |   slice : {shape : _} -> MultiSlice shape -> Shape
 555 |   slice {shape} [] = shape
 556 |   slice {shape = (_ :: _)} (Slice {size} _ _ :: xs) = size :: slice xs
 557 |   slice {shape = (_ :: _)} (Index _ :: xs) = slice xs
 558 |   slice {shape = (_ :: _)} (DynamicSlice _ size :: xs) = size :: slice xs
 559 |   slice {shape = (_ :: _)} (DynamicIndex _ :: xs) = slice xs
 560 |
 561 | ||| Slice or index `Tensor` axes. Each axis can be sliced or indexed, and this can be done with
 562 | ||| either static (`Nat`) or dynamic (scalar `U64`) indices.
 563 | |||
 564 | ||| **Static indices**
 565 | |||
 566 | ||| Static indices are `Nat`s. For example, for
 567 | ||| ```
 568 | ||| x : Tensor [5, 6] S32
 569 | ||| x = tensor [[ 0,  1,  2,  3,  4,  5],
 570 | |||             [ 6,  7,  8,  9, 10, 11],
 571 | |||             [12, 13, 14, 15, 16, 17],
 572 | |||             [18, 19, 20, 21, 22, 23],
 573 | |||             [24, 25, 26, 27, 28, 29]]
 574 | ||| ```
 575 | ||| we can index as `slice [at 1] x` to get
 576 | ||| ```
 577 | ||| x : Tensor [6] S32
 578 | ||| x = tensor [6, 7, 8, 9, 10, 11]
 579 | ||| ```
 580 | ||| or we can slice as `slice [2.to 4] x` to get
 581 | ||| ```
 582 | ||| x : Tensor [2, 6] S32
 583 | ||| x = tensor [[12, 13, 14, 15, 16, 17],
 584 | |||             [18, 19, 20, 21, 22, 23]]
 585 | ||| ```
 586 | ||| Note that in `2.to 4`, the 2 is inclusive, and the 4 exclusive, so we return indices 2 and 3.
 587 | |||
 588 | ||| **Dynamic indices**
 589 | |||
 590 | ||| Dynamic indices are scalar `U64` values, and the API works slightly differently because we
 591 | ||| can't know the value of dynamic indices until the graph is executed. For indexing, with scalar
 592 | ||| `U64` index `i` in `slice [at i] x`, `i` is clamped to be a valid index into that dimension.
 593 | ||| For example, for `i = tensor 1`, `slice [at i] x` is
 594 | ||| ```
 595 | ||| x : Tensor [6] S32
 596 | ||| x = tensor [6, 7, 8, 9, 10, 11]
 597 | ||| ```
 598 | ||| as in the static case. However, for `i = tensor 10`, `slice [at i] x` returns the last row
 599 | ||| ```
 600 | ||| x : Tensor [6] S32
 601 | ||| x = tensor [24, 25, 26, 27, 28, 29]
 602 | ||| ```
 603 | ||| We can also slice by specifying a scalar `U64` start index, and a static size, as
 604 | ||| `slice [i.size 2] x` with `i = tensor 2` to get
 605 | ||| ```
 606 | ||| x : Tensor [2, 6] S32
 607 | ||| x = tensor [[12, 13, 14, 15, 16, 17],
 608 | |||             [18, 19, 20, 21, 22, 23]]
 609 | ||| ```
 610 | ||| For a given slice `size`, the dynamic start index is clamped such that we always get `size`
 611 | ||| elements along that axis. For example, `slice [i.size 2] x` with `i = tensor 4` is
 612 | ||| ```
 613 | ||| x : Tensor [2, 6] S32
 614 | ||| x = tensor [[18, 19, 20, 21, 22, 23],
 615 | |||             [24, 25, 26, 27, 28, 29]]
 616 | ||| ```
 617 | ||| which starts at index 3 rather than index 4.
 618 | |||
 619 | ||| **Mixed static, dynamic, slicing and indexing**
 620 | |||
 621 | ||| Each axis can only be sliced or indexed, and must use only static or dynamic indices. However,
 622 | ||| across axes, we can mix these four arbitrarily. For example, with `slice [2.to 4, at 1] x` to
 623 | ||| get
 624 | ||| ```
 625 | ||| x : Tensor [2] S32
 626 | ||| x = tensor [13, 19]
 627 | ||| ```
 628 | ||| or with `i = tensor 2` in `slice [at 1, i.size 2] x` to get
 629 | ||| ```
 630 | ||| x : Tensor [2] S32
 631 | ||| x = tensor [7, 8]
 632 | ||| ```
 633 | |||
 634 | ||| Slices and indices apply to the leading axes of the tensor. For trailing axes omitted from the
 635 | ||| multi-dimensional slice, the whole of the axis is returned. If we want to slice or index over
 636 | ||| later axes and retain all indices in a leading axis, we can use the convenience function `all`,
 637 | ||| as `slice [all, at 3] x` to get
 638 | ||| ```
 639 | ||| x : Tensor [5] S32
 640 | ||| x = tensor [[3], [9], [15], [21], [27]]
 641 | ||| ```
 642 | ||| This is exactly the same as the more manual `slice [0.to 5, at 3] x` and
 643 | ||| `slice [(tensor 0).size 5, at 3] x`.
 644 | |||
 645 | ||| @at The multi-dimensional slices and indices at which to slice the tensor.
 646 | export
 647 | slice : (at : MultiSlice shape) -> Tensor shape dtype -> Tensor (slice at) dtype
 648 | slice at $ MkTensor x = MkTensor $
 649 |   let x = val0 $ Slice (mapd start (const 0) at) (mapd stop id at) (replicate (length shape) 1) x
 650 |       -- we shortcut DynamicSlice to allow autodiff for static slicing
 651 |       x = if isDynamic at then val0 $ DynamicSlice (dynStarts [] at) (mapd size id at) x else x
 652 |    in val0 $ Reshape dtype (MultiSlice.slice at) x
 653 |
 654 |       where
 655 |       mapd : ((Nat -> a) -> {d : Nat} -> SliceOrIndex d -> a) ->
 656 |              (Nat -> a) ->
 657 |              {shape : Shape} ->
 658 |              MultiSlice shape ->
 659 |              List a
 660 |       mapd _ dflt {shape} [] = Prelude.map dflt shape
 661 |       mapd f dflt (x :: xs) = f dflt x :: mapd f dflt xs
 662 |
 663 |       start : (Nat -> Nat) -> {d : Nat} -> SliceOrIndex d -> Nat
 664 |       start _ (Slice from _) = from
 665 |       start _ (Index idx) = idx
 666 |       start f {d} _ = f d
 667 |
 668 |       stop : (Nat -> Nat) -> {d : Nat} -> SliceOrIndex d -> Nat
 669 |       stop _ (Slice _ to) = to
 670 |       stop _ (Index idx) = S idx
 671 |       stop f {d} _ = f d
 672 |
 673 |       size : (Nat -> Nat) -> {d : Nat} -> SliceOrIndex d -> Nat
 674 |       size _ (Slice {size = size'} _ _) = size'
 675 |       size _ (Index _) = 1
 676 |       size _ (DynamicSlice _ size') = size'
 677 |       size _ (DynamicIndex _) = 1
 678 |
 679 |       zero : Value
 680 |       zero = val0 $ Lit [] U64 $ Scalar 0
 681 |
 682 |       isDynamic : {shape : _} -> MultiSlice shape -> Bool
 683 |       isDynamic [] = False
 684 |       isDynamic {shape = (_ :: _)} (DynamicSlice _ _ :: _) = True
 685 |       isDynamic {shape = (_ :: _)} (DynamicIndex _ :: _) = True
 686 |       isDynamic (_ :: ds) = isDynamic ds
 687 |
 688 |       dynStarts : List Value -> {shape : _} -> MultiSlice shape -> List Value
 689 |       dynStarts idxs {shape} [] = replicate (length shape) zero ++ idxs
 690 |       dynStarts idxs (DynamicSlice (MkTensor i) _ :: ds) = i :: dynStarts idxs ds
 691 |       dynStarts idxs (DynamicIndex (MkTensor i) :: ds) = i :: dynStarts idxs ds
 692 |       dynStarts idxs (_ :: ds) = zero :: dynStarts idxs ds
 693 |
 694 | ||| Concatenate two `Tensor`s along the specified `axis`. For example,
 695 | ||| `concat 0 (tensor [[1, 2], [3, 4]]) (tensor [[5, 6]])` and
 696 | ||| `concat 1 (tensor [[3], [6]]) (tensor [[4, 5], [7, 8]])` are both
 697 | ||| `tensor [[1, 2], [3, 4], [5, 6]]`.
 698 | export
 699 | concat :
 700 |   (axis : Nat) ->
 701 |   Tensor s dtype ->
 702 |   Tensor s' dtype ->
 703 |   {auto 0 inBounds : (InBounds axis s, InBounds axis s')} ->
 704 |   {auto 0 shapesConcatenable : deleteAt axis s = deleteAt axis s'} ->
 705 |   Tensor (replaceAt axis (index axis s + index axis s') s) dtype
 706 | concat axis (MkTensor x) (MkTensor x') = t0 $ Concat axis [x, x']
 707 |
 708 | ||| Transpose a matrix. For example, `(tensor [[1, 2], [3, 4]]).T` is `tensor [[1, 3], [2, 4]]`.
 709 | export
 710 | (.T) : Tensor [m, n] dtype -> Tensor [n, m] dtype
 711 | (MkTensor x).T = t0 $ Transpose [1, 0] x
 712 |
 713 | ||| Transpose axes of a tensor. This is a more general version of `(.T)`, in which you can
 714 | ||| transpose any number of axes in a tensor of arbitrary rank. The i'th axis in the resulting
 715 | ||| tensor corresponds to the `index i ordering`'th axis in the input tensor. For example, for
 716 | ||| ```
 717 | ||| x : Tensor [2, 3, 4] S32
 718 | ||| x = tensor [[[ 0,  1,  2,  3],
 719 | |||              [ 4,  5,  6,  7],
 720 | |||              [ 8,  9, 10, 11]],
 721 | |||             [[12, 13, 14, 15],
 722 | |||              [16, 17, 18, 19],
 723 | |||              [20, 21, 22, 23]]]
 724 | ||| ```
 725 | ||| `transpose [0, 2, 1] x` is
 726 | ||| ```
 727 | ||| x : Tensor [2, 4, 3] S32
 728 | ||| x = tensor [[[ 0,  4,  8],
 729 | |||              [ 1,  5,  9],
 730 | |||              [ 2,  6, 10],
 731 | |||              [ 3,  7, 11]],
 732 | |||             [[12, 16, 20],
 733 | |||              [13, 17, 21],
 734 | |||              [14, 18, 22],
 735 | |||              [15, 19, 23]]]
 736 | ||| ```
 737 | ||| `transpose [2, 0, 1] x` is
 738 | ||| ```
 739 | ||| x : Tensor [4, 2, 3] S32
 740 | ||| x = tensor [[[ 0,  4,  8],
 741 | |||              [12, 16, 20]],
 742 | |||             [[ 1,  5,  9],
 743 | |||              [13, 17, 21]],
 744 | |||             [[ 2,  6, 10],
 745 | |||              [14, 18, 22]],
 746 | |||             [[ 3,  7, 11],
 747 | |||              [15, 19, 23]]]
 748 | ||| ```
 749 | |||
 750 | ||| In order to see what effect transposing a tensor has, it can help to bear in mind the following:
 751 | ||| * if an element can be found with `slice [at 3, at 4, at 5] x` in the original tensor,
 752 | |||   that same element can instead be found with `slice [at 5, at 3, at 4]` given a
 753 | |||   `transpose [2, 0, 1]`. That is, transposing axes re-orders indices when indexing.
 754 | ||| * with `transpose [2, 0, 1]`, traversing the first axis in the result is equivalent to
 755 | |||   traversing the last axis in the input. Similarly, traversing the last axis in the result is
 756 | |||   equivalent to traversing the second axis in the input.
 757 | export
 758 | transpose :
 759 |   (ordering : List Nat) ->
 760 |   Tensor shape dtype ->
 761 |   {auto 0 lengths : length ordering = length shape} ->
 762 |   {auto 0 axesUnique : unique ordering = True} ->
 763 |   {auto 0 inBounds : All (flip InBounds shape) ordering} ->
 764 |   Tensor (multiIndex ordering shape) dtype
 765 | transpose ordering $ MkTensor x = t0 $ Transpose ordering x
 766 |
 767 | ||| A `DimBroadcastable from to` proves that a dimension of size `from` can be broadcast to a
 768 | ||| dimension of size `to`.
 769 | public export
 770 | data DimBroadcastable : (0 from : Nat) -> (0 to : Nat) -> Type where
 771 |   ||| Proof that any dimension can be broadcast to itself. For example in shapes `[2, 3]` to
 772 |   ||| `[2, 3]`.
 773 |   Same : DimBroadcastable x x
 774 |
 775 |   ||| Proof that a dimension of length one can be broadcast to any size. For example in shapes
 776 |   ||| `[2, 1]` to `[2, 3]`
 777 |   Stack : DimBroadcastable 1 _
 778 |
 779 |   ||| Proof that any dimension can be broadcast to zero. For example in shapes `[2, 3]` to `[2, 0]`.
 780 |   Zero : DimBroadcastable _ 0
 781 |
 782 | namespace Broadcastable
 783 |   ||| A `Broadcastable from to` constitutes proof that the shape `from` can be broadcast to the
 784 |   ||| shape `to`.
 785 |   public export
 786 |   data Broadcastable : (0 from : Shape) -> (0 to : Shape) -> Type where
 787 |     ||| Proof that a shape can be broadcast to itself. For example:
 788 |     |||
 789 |     ||| [] to []
 790 |     ||| [3, 4] to [3, 4]
 791 |     |||
 792 |     ||| Implementation note: we could have used `Broadcast [] []`, which would have resulted in more
 793 |     ||| atomic constructors for `Broadcastable`, but the author guesses that this implementation helps
 794 |     ||| the type checker avoid applications of `Match`.
 795 |     Same : Broadcastable x x
 796 |
 797 |     ||| Proof that a dimension of size `f` can be broadcast to size `t` if these dimensions
 798 |     ||| are `DimBroadcastable f t`. For example:
 799 |     |||
 800 |     ||| [2, 3] to [2, 3]
 801 |     ||| [2, 1] to [2, 3]
 802 |     ||| [2, 1] to [2, 0]
 803 |     Match : forall from, to .
 804 |             {auto 0 ranksEq : length from = length to} ->
 805 |             {auto 0 dimBroadcastable : DimBroadcastable f t} ->
 806 |             Broadcastable from to ->
 807 |             Broadcastable (f :: from) (t :: to)
 808 |
 809 |     ||| Proof that broadcasting can add outer dimensions i.e. nesting. For example:
 810 |     |||
 811 |     ||| [3] to [1, 3]
 812 |     ||| [3] to [5, 3]
 813 |     Nest : Broadcastable f t -> Broadcastable f (_ :: t)
 814 |
 815 | ||| A shape can be extended with any number of leading dimensions.
 816 | |||
 817 | ||| @leading The leading dimensions.
 818 | export
 819 | broadcastableByLeading : (leading : List Nat) -> Broadcastable shape (leading ++ shape)
 820 | broadcastableByLeading [] = Same
 821 | broadcastableByLeading (l :: ls) = Nest (broadcastableByLeading ls)
 822 |
 823 | ||| A scalar can be broadcast to any shape.
 824 | %hint
 825 | export
 826 | scalarToAnyOk : (to : Shape) -> Broadcastable [] to
 827 | scalarToAnyOk to = rewrite sym $ appendNilRightNeutral to in broadcastableByLeading to
 828 |
 829 | ||| Broadcast a `Tensor` to a new compatible shape. For example,
 830 | ||| ```
 831 | ||| x : Tensor [2, 3] S32
 832 | ||| x = broadcast (tensor [4, 5, 6])
 833 | ||| ```
 834 | ||| is
 835 | ||| ```
 836 | ||| x : Tensor [2, 3] S32
 837 | ||| x = tensor [[4, 5, 6], [4, 5, 6]]
 838 | ||| ```
 839 | export
 840 | broadcast :
 841 |   {to : _} -> {dtype : _} ->
 842 |   {auto shapesOK : Broadcastable from to} ->
 843 |   Tensor from dtype ->
 844 |   Tensor to dtype
 845 | broadcast $ MkTensor {shape = _} x = t0 $ Broadcast (Explicit to) x
 846 |
 847 | ||| A `Tensor` where every element has the specified value. For example,
 848 | ||| ```
 849 | ||| fives : Tensor [2, 3] S32
 850 | ||| fives = fill 5
 851 | ||| ```
 852 | ||| is
 853 | ||| ```
 854 | ||| fives : Tensor [2, 3] S32
 855 | ||| fives = tensor [[5, 5, 5],
 856 | |||                 [5, 5, 5]]
 857 | ||| ```
 858 | export
 859 | fill : {shape : _} -> {dtype : _} -> idrisType dtype -> Tensor shape dtype
 860 | fill x = broadcast {shapesOK = scalarToAnyOk shape} (tensor (Scalar x))
 861 |
 862 | ||| A constant where values increment from zero along the specified `axis`. For example,
 863 | ||| ```
 864 | ||| x : Tensor [3, 5] S32
 865 | ||| x = iota 1
 866 | ||| ```
 867 | ||| is the same as
 868 | ||| ```
 869 | ||| x : Tensor [3, 5] S32
 870 | ||| x = tensor [[0, 1, 2, 3, 4],
 871 | |||             [0, 1, 2, 3, 4],
 872 | |||             [0, 1, 2, 3, 4]]
 873 | ||| ```
 874 | ||| and
 875 | ||| ```
 876 | ||| x : Tensor [3, 5] S32
 877 | ||| x = iota 0
 878 | ||| ```
 879 | ||| is the same as
 880 | ||| ```
 881 | ||| x : Tensor [3, 5] S32
 882 | ||| x = tensor [[0, 0, 0, 0, 0],
 883 | |||             [1, 1, 1, 1, 1],
 884 | |||             [2, 2, 2, 2, 2]]
 885 | ||| ```
 886 | export
 887 | iota :
 888 |   {shape : _} ->
 889 |   {dtype : _} ->
 890 |   {auto 0 _ : Num dtype} ->
 891 |   (axis : Nat) ->
 892 |   {auto 0 inBounds : InBounds axis shape} ->
 893 |   Tensor shape dtype
 894 | iota dimension = t0 $ Iota shape dtype dimension
 895 |
 896 | ||| A while-loop operating on a single tensor.
 897 | |||
 898 | ||| `while1` iteratively checks if the tensor satisfies `condition`, and if it does, updates it with
 899 | ||| `body`.
 900 | |||
 901 | ||| **Note:** For the XLA compiler, there are scoping restrictions on functions passed to `while1`.
 902 | ||| See the tutorial _Nuisances in the Tensor API_ for details.
 903 | |||
 904 | ||| @condition The guard condition for each iteration.
 905 | ||| @body The update step.
 906 | ||| @initial The initial tensor.
 907 | export covering
 908 | while1 :
 909 |   (condition : Tensor shape dtype -> Tag $ Tensor [] PRED) ->
 910 |   (body : Tensor shape dtype -> Tag $ Tensor shape dtype) ->
 911 |   (initial : Tensor shape dtype) ->
 912 |   Tag $ Tensor shape dtype
 913 | while1 condition body (MkTensor i0) =
 914 |   pure $ t0 $ While !(mkFn1 [_] [_] condition) !(mkFn1 [_] [_] body) [i0]
 915 |
 916 | ||| A while-loop operating on two tensors.
 917 | |||
 918 | ||| `while2` iteratively checks if the tensors together satisfy `condition`, and if so, updates them
 919 | ||| with `body`.
 920 | |||
 921 | ||| **Note:** For the XLA compiler, there are scoping restrictions on functions passed to `while2`.
 922 | ||| See the tutorial _Nuisances in the Tensor API_ for details.
 923 | |||
 924 | ||| @condition The guard condition for each iteration.
 925 | ||| @body The update step.
 926 | ||| @initial One initial tensor.
 927 | ||| @initial' The other initial tensor.
 928 | export covering
 929 | while2 :
 930 |   (condition : Tensor s a -> Tensor s' a' -> Tag $ Tensor [] PRED) ->
 931 |   (body : Tensor s a -> Tensor s' a' -> Tag $ All2 Tensor [s, s'] [a, a']) ->
 932 |   (initial : Tensor s a) -> (initial' : Tensor s' a') ->
 933 |   Tag $ All2 Tensor [s, s'] [a, a']
 934 | while2 condition body (MkTensor i) (MkTensor i') = do
 935 |   res <- tag $ Concrete $ While !(mkFn1 [_, _] [_, _] condition) !(mkFn [_, _] [_, _] body) [i, i']
 936 |   pure [MkTensor $ V 0 res, MkTensor $ V 1 res]
 937 |
 938 | ||| Lift a unary function on scalars to an element-wise function on `Tensor`s of arbitrary shape.
 939 | ||| For example,
 940 | ||| ```
 941 | ||| recip : Tensor [] F64 -> Tag $ Tensor [] F64
 942 | ||| recip x = pure $ 1.0 / x
 943 | ||| ```
 944 | ||| can be lifted to an element-wise reciprocal function as `map recip (tensor [-2, 0.4])`,
 945 | ||| which produces `tensor [-0.5, 2.5]`.
 946 | |||
 947 | ||| **Note:** For the XLA compiler, there are scoping restrictions on the function passed to `map`.
 948 | ||| See the tutorial _Nuisances in the Tensor API_ for details.
 949 | export
 950 | map : {b : _} -> (Tensor [] a -> Tag $ Tensor [] b) -> Tensor shape a -> Tag $ Tensor shape b
 951 | map f $ MkTensor {shape = _} x =
 952 |   pure $ t0 $ Map !(mkFn1 [_] [_] f) [x] (TensorType shape b) (range $ length shape)
 953 |
 954 | ||| Lift a binary function on scalars to an element-wise function on `Tensor`s of arbitrary shape.
 955 | ||| For example,
 956 | ||| ```
 957 | ||| addRecip : Tensor [] F64 -> Tensor [] F64 -> Tag $ Tensor [] F64
 958 | ||| addRecip x y = pure $ x + 1.0 / y
 959 | ||| ```
 960 | ||| can be lifted to an element-wise function as
 961 | ||| `map2 addRecip (tensor [3.0, -3.0]) (tensor [-2.0, 0.4])`, which produces `tensor [2.5, -0.5]`.
 962 | |||
 963 | ||| **Note:** For the XLA compiler, there are scoping restrictions on the function passed to `map2`.
 964 | ||| See the tutorial _Nuisances in the Tensor API_ for details.
 965 | export
 966 | map2 :
 967 |   {c : _} ->
 968 |   (Tensor [] a -> Tensor [] b -> Tag $ Tensor [] c) ->
 969 |   Tensor shape a -> Tensor shape b -> Tag $ Tensor shape c
 970 | map2 f (MkTensor {shape = _} x) (MkTensor x') =
 971 |   pure $ t0 $ Map !(mkFn1 [_, _] [_, _] f) [x, x'] (TensorType shape c) (range $ length shape)
 972 |
 973 | ||| Reduce elements along one `axis` of a `Tensor` according to a specified `reducer` `Monoid`.
 974 | ||| For example, if `x = tensor [[0, 1, 2], [3, 4, 5]]`, then reduce @{Sum} 0 x` produces
 975 | ||| `tensor [3, 5, 7]`, and `reduce @{Sum} 1 x` produces `tensor [3, 12]`.
 976 | |||
 977 | ||| **Note:** `Semigroup` doesn't use `Tag`, which limits the functions that can be used in
 978 | ||| `reduce`. However, the most commonly used semigroups don't need `Tag`, including `Sum`,
 979 | ||| `Prod`, `Min` and `Max`, so for ergonomics, we have opted to use `Monoid` as is. We can
 980 | ||| provide an overloaded variant if requested.
 981 | |||
 982 | ||| **Note:** For the XLA compiler, there are scoping restrictions on the monoid's semigroup.
 983 | ||| See the tutorial _Nuisances in the Tensor API_ for details.
 984 | |||
 985 | ||| @reducer How to reduce elements along the given `axis`.
 986 | ||| @axis The axis along which to reduce elements.
 987 | export
 988 | reduce :
 989 |   (reducer : Monoid (Tensor [] dtype)) =>
 990 |   (axes : List Nat) ->
 991 |   {auto 0 axesUnique : Sorted LT axes} ->
 992 |   {auto 0 axesInBounds : All (flip InBounds shape) axes} ->
 993 |   Tensor shape dtype ->
 994 |   Tag $ Tensor (deleteAt axes shape) dtype
 995 | reduce axes $ MkTensor x = do
 996 |   let semigroup : Monoid a -> Semigroup a
 997 |       semigroup _ = %search
 998 |
 999 |   g <- mkFn1 [_, _] [_, _] (pure .: (<+>) @{semigroup reducer})
1000 |   let MkTensor neutral' = neutral @{reducer}
1001 |   pure $ t0 $ Reduce g [neutral'] axes [x]
1002 |
1003 | ||| Sort the elements of a `Tensor` along a specified `dimension` according to a scalar-wise
1004 | ||| ordering. For sorting function `f`, elements are sorted such that for consecutive sorted
1005 | ||| elements `a` and `b`, either `f a b` is true, or `f a b` *and* `f b a` are false.
1006 | |||
1007 | ||| **Note:** Sorting is not stable, meaning elements that compare equal according the ordering may
1008 | ||| be sorted in a different order to the order they appear in the input.
1009 | |||
1010 | ||| **Note:** `sort` is limited to use comparison function without `Tag`. However, since the most
1011 | ||| commonly-used functions, including (>), (<), (>=), and (<=), don't use `Tag`, we have opted to
1012 | ||| omit it for ergonomics. We can trivially provide an overloaded variant if requested.
1013 | |||
1014 | ||| For example, for `x = tensor [[1, 6, 4], [3, 2, 5]]`, `sort (<) 0 x` produces
1015 | ||| `tensor [[1, 2, 4], [3, 6, 5]]`, while `sort (<) 1 x` produces
1016 | ||| `tensor [[1, 4, 6], [2, 3, 5]]`.
1017 | |||
1018 | ||| **Note:** For the XLA compiler, there are scoping restrictions on the function passed to `sort`.
1019 | ||| See the tutorial _Nuisances in the Tensor API_ for details.
1020 | export
1021 | sort :
1022 |   (Tensor [] dtype -> Tensor [] dtype -> Tensor [] PRED) ->
1023 |   (dimension : Nat) ->
1024 |   Tensor shape dtype ->
1025 |   {auto 0 dimInBounds : InBounds dimension shape} ->
1026 |   Tag $ Tensor shape dtype
1027 | sort comp dimension $ MkTensor x =
1028 |   pure $ t0 $ Sort !(mkFn1 [_, _] [_, _] $ pure .: comp) dimension False x
1029 |
1030 | ||| Reverse elements along the specified axes. For example, for
1031 | ||| ```
1032 | ||| x : Tensor [2, 3] S32
1033 | ||| x = tensor [[-2, -1,  0],
1034 | |||             [ 1,  2,  3]]
1035 | ||| ```
1036 | ||| `reverse [0] x` is
1037 | ||| ```
1038 | ||| x : Tensor [2, 3] S32
1039 | ||| x = tensor [[ 1,  2,  3],
1040 | |||             [-2, -1,  0]]
1041 | ||| ```
1042 | ||| `reverse [1] x` is
1043 | ||| ```
1044 | ||| x : Tensor [2, 3] S32
1045 | ||| x = tensor [[ 0, -1, -2],
1046 | |||             [ 3,  2,  1]]
1047 | ||| ```
1048 | ||| and `reverse [0, 1] x` is
1049 | ||| ```
1050 | ||| x : Tensor [2, 3] S32
1051 | ||| x = tensor [[ 3,  2,  1],
1052 | |||             [ 0, -1, -2]]
1053 | ||| ```
1054 | |||
1055 | ||| **Note:** This function requires `axes` is ordered simply so that elements are unique.
1056 | ||| The ordering itself is irrelevant to the implementation, but ensures uniqueness without using
1057 | ||| proofs of contradiction that can be difficult for Idris to construct.
1058 | export
1059 | reverse :
1060 |   (axes : List Nat) ->
1061 |   {auto 0 axesUnique : Sorted LT axes} ->
1062 |   {auto 0 axesInBounds : All (flip InBounds shape) axes} ->
1063 |   Tensor shape dtype ->
1064 |   Tensor shape dtype
1065 | reverse axes $ MkTensor x = t0 $ Reverse axes x
1066 |
1067 | ewUnary : UnaryOp -> Tensor s a -> Tensor s a
1068 | ewUnary op $ MkTensor x = t0 $ UnaryElementwise op x
1069 |
1070 | ewBinary : BinaryOp -> Tensor s a -> Tensor s a -> Tensor s a
1071 | ewBinary op (MkTensor x) (MkTensor x') = t0 $ BinaryElementwise op x x'
1072 |
1073 | ewBinary' : {out : _} -> BinaryOp -> Tensor s a -> Tensor s a -> Tensor s out
1074 | ewBinary' op (MkTensor x) (MkTensor x') = t0 $ BinaryElementwise op x x'
1075 |
1076 | ||| Element-wise equality. For example, `tensor [1, 2] /= tensor [1, 3]` is
1077 | ||| `tensor [True, False]`.
1078 | export
1079 | (==) : Eq dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape PRED
1080 | (==) = ewBinary' $ Compare Eq
1081 |
1082 | ||| Element-wise inequality. For example, `tensor [1, 2] /= tensor [1, 3]` is
1083 | ||| `tensor [False, True]`.
1084 | export
1085 | (/=) : Eq dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape PRED
1086 | (/=) = ewBinary' $ Compare Ne
1087 |
1088 | ||| Element-wise less than. For example, `tensor [1, 2, 3] < tensor [2, 2, 2]` is
1089 | ||| `tensor [True, False, False]`.
1090 | export
1091 | (<) : Ord dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape PRED
1092 | (<) = ewBinary' $ Compare Lt
1093 |
1094 | ||| Element-wise greater than. For example, `tensor [1, 2, 3] > tensor [2, 2, 2]` is
1095 | ||| `tensor [False, False, True]`.
1096 | export
1097 | (>) : Ord dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape PRED
1098 | (>) = ewBinary' $ Compare Gt
1099 |
1100 | ||| Element-wise less than or equal. For example, `tensor [1, 2, 3] <= tensor [2, 2, 2]`
1101 | ||| is `tensor [True, True, False]`.
1102 | export
1103 | (<=) : Ord dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape PRED
1104 | (<=) = ewBinary' $ Compare Le
1105 |
1106 | ||| Element-wise greater than or equal. For example,
1107 | ||| `tensor [1, 2, 3] >= tensor [2, 2, 2]` is `tensor [False, True, True]`.
1108 | export
1109 | (>=) : Ord dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape PRED
1110 | (>=) = ewBinary' $ Compare Ge
1111 |
1112 | ||| Element-wise boolean and. For example,
1113 | ||| `tensor [True, True, False, False] && tensor [True, False, True, False]` is
1114 | ||| `tensor [True, False, False, False]`.
1115 | export
1116 | (&&) : Tensor shape PRED -> Tensor shape PRED -> Tensor shape PRED
1117 | (&&) = ewBinary And
1118 |
1119 | namespace Semigroup
1120 |   export
1121 |   [All] Semigroup (Tensor shape PRED) where
1122 |     (<+>) = (&&)
1123 |
1124 | namespace Monoid
1125 |   export
1126 |   [All] {shape : _} -> Monoid (Tensor shape PRED) using Tensor.Semigroup.All where
1127 |     neutral = fill True
1128 |
1129 | ||| Element-wise boolean or. For example,
1130 | ||| `tensor [True, True, False, False] || tensor [True, False, True, False]` is
1131 | ||| `tensor [True, True, True, False]`.
1132 | export
1133 | (||) : Tensor shape PRED -> Tensor shape PRED -> Tensor shape PRED
1134 | (||) = ewBinary Or
1135 |
1136 | namespace Semigroup
1137 |   export
1138 |   [Any] Semigroup (Tensor shape PRED) where
1139 |     (<+>) = (||)
1140 |
1141 | namespace Monoid
1142 |   export
1143 |   [Any] {shape : _} -> Monoid (Tensor shape PRED) using Tensor.Semigroup.Any where
1144 |     neutral = fill False
1145 |
1146 | ||| Element-wise boolean negation. For example, `not (tensor [True, False])` is
1147 | ||| `tensor [False, True]`.
1148 | export
1149 | not : Tensor shape PRED -> Tensor shape PRED
1150 | not = ewUnary Not
1151 |
1152 | ||| Choose elements from two `Tensor`s based on a `Tensor` of predicates. For each element in the
1153 | ||| predicates, the output will use the corresponding element from `onTrue` if the element is
1154 | ||| truthy, else the element from `onFalse`. For example, for
1155 | ||| ```
1156 | ||| preds : Tensor [3] PRED
1157 | ||| preds = tensor [False, True, False]
1158 | |||
1159 | ||| onTrue : Tensor [3] S32
1160 | ||| onTrue = tensor [1, 2, 3]
1161 | |||
1162 | ||| onFalse : Tensor [3] S32
1163 | ||| onFalse = tensor [4, 5, 6]
1164 | ||| ```
1165 | ||| `select preds onTrue onFalse` is `tensor [4, 2, 6]`.
1166 | |||
1167 | ||| @onTrue The elements to choose where the predicate elements are truthy.
1168 | ||| @onFalse The elements to choose where the predicate elements are falsy.
1169 | export
1170 | select :
1171 |   Tensor shape PRED ->
1172 |   (onTrue, onFalse : Tensor shape dtype) ->
1173 |   Tensor shape dtype
1174 | select (MkTensor p) (MkTensor t) (MkTensor f) = t0 $ Select p t f
1175 |
1176 | ||| Use a scalar predicate to evaluate one of two branches. If the predicate is truthy,
1177 | ||| evaluate `onTrue`, else `onFalse`. Each branch is evaluated lazily; only one will be
1178 | ||| evaluated.
1179 | |||
1180 | ||| For example, for
1181 | ||| ```
1182 | ||| f : Tensor [] F64 -> Tag $ Tensor [] F64
1183 | ||| f x = if_ (x < 1.0) (pure $ cos x) (do x <- tag x; x * x)
1184 | ||| ```
1185 | ||| `f 0.0` produces `1.0`, and `f 4.0` produces `8.0`.
1186 | |||
1187 | ||| **Note:** For the XLA compiler, there are scoping restrictions on branches used in `if_`, as
1188 | ||| StableHLO gives branches their own scope. See the tutorial _Nuisances in the Tensor API_
1189 | ||| for details.
1190 | |||
1191 | ||| @onTrue The branch to evaluate if the predicate is truthy.
1192 | ||| @onFalse The branch to evaluate if the predicate is falsy.
1193 | export
1194 | if_ :
1195 |   {shape : _} -> {dtype : _} ->
1196 |   Tensor [] PRED ->
1197 |   (onTrue, onFalse : Tag $ Tensor shape dtype) ->
1198 |   Tag $ Tensor shape dtype
1199 | if_ (MkTensor pred) onTrue onFalse =
1200 |   pure $ t0 $ If (TensorType shape dtype) pred !(mkFn1 [] [] onTrue) !(mkFn1 [] [] onFalse)
1201 |
1202 | ||| The identity tensor, with inferred shape and element type. For example,
1203 | ||| ```
1204 | ||| x : Tensor [2, 2] S32
1205 | ||| x = identity
1206 | ||| ```
1207 | ||| is
1208 | ||| ```
1209 | ||| x : Tensor [2, 2] S32
1210 | ||| x = tensor [[1, 0],
1211 | |||             [0, 1]]
1212 | ||| ```
1213 | export
1214 | identity : {n : _} -> {dtype : _} -> Num dtype => Tensor [n, n] dtype
1215 | identity =
1216 |   let MkTensor x = iota 0 {shape = [n, n], dtype = U64} == iota 1
1217 |    in t0 $ Convert dtype [n, n] x
1218 |
1219 | -- see https://www.python.org/dev/peps/pep-0465/#precedence-and-associativity
1220 | export infixl 9 @@
1221 |
1222 | namespace Vector
1223 |   ||| Vector dot product with a tensor of any rank. The vector dot product is with the first axis of
1224 |   ||| the right-hand side tensor. For example `tensor [0, 1, 2] @@ tensor [-1, -3, -1]` is
1225 |   ||| `-1`.
1226 |   export
1227 |   (@@) : Num dtype => Tensor [S m] dtype -> Tensor [S m] dtype -> Tensor [] dtype
1228 |   (MkTensor x) @@ (MkTensor x') = t0 $ DotGeneral [] [] [0] [0] (TensorType [] dtype) x x'
1229 |
1230 | namespace Matrix
1231 |   ||| Matrix multiplication with a matrix or vector. Contraction is along the last axis of the first
1232 |   ||| and the first axis of the last. For example,
1233 |   ||| ```
1234 |   ||| x : Tensor [2, 3] S32
1235 |   ||| x = tensor [[-1, -2, -3],
1236 |   |||             [ 0,  1,  2]]
1237 |   |||
1238 |   ||| y : Tensor [3, 1] S32
1239 |   ||| y = tensor [[4, 0, 5]]
1240 |   |||
1241 |   ||| z : Tensor [2, 1] S32
1242 |   ||| z = x @@ y
1243 |   ||| ```
1244 |   ||| is
1245 |   ||| ```
1246 |   ||| z : Tensor [2, 1] S32
1247 |   ||| z = tensor [-19, 10]
1248 |   ||| ```
1249 |   export
1250 |   (@@) : Num dtype =>
1251 |          Tensor [n, S m] dtype ->
1252 |          Tensor (S m :: tl) dtype ->
1253 |          {auto 0 vectorTail : length tl `LTE` 1} ->
1254 |          Tensor (n :: tl) dtype
1255 |   (MkTensor x) @@ (MkTensor x') = t0 $ DotGeneral [] [] [1] [0] (TensorType (n :: tl) dtype) x x'
1256 |
1257 | ||| The output shape of a `dotGeneral` operation.
1258 | public export
1259 | contract : (lBatch, rBatch, lContract, rContract : List Nat) ->
1260 |            (ls, rs : Shape) ->
1261 |            {auto 0 lInBoundsBatch : All (flip InBounds ls) lBatch} ->
1262 |            {auto 0 rInBoundsBatch : All (flip InBounds rs) rBatch} ->
1263 |            {auto 0 lInBoundsContract : All (flip InBounds ls) lContract} ->
1264 |            {auto 0 rInBoundsContract : All (flip InBounds rs) rContract} ->
1265 |            Shape
1266 | contract lBatch rBatch lContract rContract ls rs =
1267 |   let lResultDims = deleteAt {inBounds = lInBoundsBatch ++ lInBoundsContract}
1268 |                              (lBatch ++ lContract) ls
1269 |       rResultDims = deleteAt {inBounds = rInBoundsBatch ++ rInBoundsContract}
1270 |                              (rBatch ++ rContract) rs
1271 |    in multiIndex lBatch ls ++ lResultDims ++ rResultDims
1272 |
1273 | ||| Matrix multiplication.
1274 | |||
1275 | ||| This is a much more general version of `(@@)`, in which you can specify any number of batch
1276 | ||| and contracting axes. Matrix multiplication is done over each contracting axis.
1277 | ||| The operation is vectorized over batch axes. For each contracting axis on the left-hand
1278 | ||| operand, there is one contracting axis on the right-hand operand. These can be different axes
1279 | ||| in each operand. The same is true for each batch axis.
1280 | |||
1281 | ||| For example, we can vectorize over a typical rank-two matrix multiplication as follows: given
1282 | ||| two inputs tensors
1283 | ||| ```
1284 | ||| let x : Tensor [3, 4, 5, 6] F64
1285 | |||     y : Tensor [3, 4, 6, 7] F64
1286 | ||| ```
1287 | ||| we do
1288 | ||| ```
1289 | ||| let z : Tensor [3, 4, 5, 7] F64 = dotGeneral [0, 1] [0, 1] [3] [2] x y
1290 | ||| ```
1291 | ||| Here, we vectorized over the first two axes `[0, 1]`, and do standard matrix multiplication
1292 | ||| over the remaining axes by specifying the axes 3 and 2 respectively as contracting axes. Notice
1293 | ||| how the batch axes appear once each at the start of the output shape, and the contracting axis
1294 | ||| disappears. Remaining axes appear in order from left to right.
1295 | |||
1296 | ||| Note this API is somewhat of a quickfix to bring general matrix multiplication to the tensor
1297 | |||   API. It is not thoroughly tested. Expect it to change in the future.
1298 | export
1299 | dotGeneral :
1300 |   Num dtype =>
1301 |   (lBatch, rBatch, lContract, rContract : List Nat) ->
1302 |   {auto 0 lUnique : unique (lBatch ++ lContract) = True} ->
1303 |   {auto 0 rUnique : unique (rBatch ++ rContract) = True} ->
1304 |   {auto 0 lInBoundsBatch : All (flip InBounds ls) lBatch} ->
1305 |   {auto 0 rInBoundsBatch : All (flip InBounds rs) rBatch} ->
1306 |   {auto 0 lInBoundsContract : All (flip InBounds ls) lContract} ->
1307 |   {auto 0 rInBoundsContract : All (flip InBounds rs) rContract} ->
1308 |   {auto 0 batchDimsEq : multiIndex lBatch ls = multiIndex rBatch rs} ->
1309 |   {auto 0 contractDimsEq : multiIndex lContract ls = multiIndex rContract rs} ->
1310 |   Tensor ls dtype ->
1311 |   Tensor rs dtype ->
1312 |   Tensor (contract lBatch rBatch lContract rContract ls rs) dtype
1313 | dotGeneral lb rb lc rc (MkTensor x) (MkTensor y) =
1314 |   let resultType = TensorType (contract lb rb lc rc ls rs) dtype
1315 |    in t0 $ DotGeneral lb rb lc rc resultType x y
1316 |
1317 | ||| Element-wise addition. For example, `tensor [1, 2] + tensor [3, 4]` is
1318 | ||| `tensor [4, 6]`.
1319 | export
1320 | (+) : Num dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape dtype
1321 | (+) = ewBinary Add
1322 |
1323 | namespace Semigroup
1324 |   export
1325 |   [Sum] Num dtype => Semigroup (Tensor shape dtype) where
1326 |     (<+>) = (+)
1327 |
1328 | namespace Monoid
1329 |   export
1330 |   [Sum] {shape : _} -> {dtype : _} -> Prelude.Num (idrisType dtype) => Num dtype =>
1331 |     Monoid (Tensor shape dtype) using Semigroup.Sum where
1332 |       neutral = fill 0
1333 |
1334 | ||| Element-wise negation. For example, `- tensor [1, -2]` is `tensor [-1, 2]`.
1335 | export
1336 | negate : Neg dtype => Tensor shape dtype -> Tensor shape dtype
1337 | negate $ MkTensor i = t0 $ UnaryElementwise Neg i
1338 |
1339 | ||| Element-wise subtraction. For example, `tensor [3, 4] - tensor [4, 2]` is
1340 | ||| `tensor [-1, 2]`.
1341 | export
1342 | (-) : Neg dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape dtype
1343 | (-) = ewBinary Sub
1344 |
1345 | ||| Element-wise multiplication. For example, `tensor [2, 3] * tensor [4, 5]` is
1346 | ||| `tensor [8, 15]`.
1347 | export
1348 | (*) : Num dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape dtype
1349 | (*) = ewBinary Mul
1350 |
1351 | namespace Scalarwise
1352 |   ||| Multiplication by a scalar. For example, `tensor 2 * tensor [3, 5]` is
1353 |   ||| `tensor [6, 10]`.
1354 |   |||
1355 |   ||| The RHS is required to be non-scalar simply to avoid ambiguities with element-wise `(*)`.
1356 |   export
1357 |   (*) : Num dtype => Tensor [] dtype -> Tensor (d :: ds) dtype -> Tensor (d :: ds) dtype
1358 |   l * r =
1359 |     let MkTensor {shape = _ :: _} _ = r
1360 |      in broadcast {shapesOK = scalarToAnyOk (d :: ds)} l * r
1361 |
1362 | namespace Semigroup
1363 |   export
1364 |   [Prod] Num dtype => Semigroup (Tensor shape dtype) where
1365 |     (<+>) = (*)
1366 |
1367 | namespace Monoid
1368 |   export
1369 |   [Prod] {shape : _} -> {dtype : _} -> Prelude.Num (idrisType dtype) => Num dtype =>
1370 |     Monoid (Tensor shape dtype) using Semigroup.Prod where
1371 |       neutral = fill 1
1372 |
1373 | ||| Element-wise floating point division. For example, `tensor [2, 3] / tensor [4, 5]` is
1374 | ||| `tensor [0.5, 0.6]`.
1375 | export
1376 | (/) : Fractional dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape dtype
1377 | (/) = ewBinary Div
1378 |
1379 | namespace Scalarwise
1380 |   ||| Floating point division by a scalar. For example, `tensor [3.4, -5.6] / tensor 2` is
1381 |   ||| `tensor [1.7, -2.8]`.
1382 |   |||
1383 |   ||| The LHS is required to be non-scalar simply to avoid ambiguities with element-wise `(/)`.
1384 |   export
1385 |   (/) : Fractional dtype => Tensor (d :: ds) dtype -> Tensor [] dtype -> Tensor (d :: ds) dtype
1386 |   l / r =
1387 |     let MkTensor {shape = _ :: _} _ = l
1388 |      in l / broadcast {shapesOK = scalarToAnyOk (d :: ds)} r
1389 |
1390 | inf = tensor 1.0 / tensor 0.0
1391 | nan = tensor 0.0 / tensor 0.0
1392 |
1393 | ||| Element-wise division of natural numbers. For example,
1394 | ||| `div (tensor [13, 8]) [3, 4]` is `tensor [4, 2]`.
1395 | export
1396 | div : Tensor shape U64 ->
1397 |       (denom : Literal shape Nat) ->
1398 |       {auto 0 isSucc : All IsSucc denom} ->
1399 |       Tensor shape U64
1400 | div x y with (x)
1401 |   _ | (MkTensor {shape = _} _) = ewBinary Div x (tensor {dtype = U64} $ cast <$> y)
1402 |
1403 | namespace Integral
1404 |   ||| Element-wise remainder for natural numbers. For example,
1405 |   ||| `rem (tensor [13, 8]) [3, 4]` is `tensor [1, 0]`.
1406 |   export
1407 |   rem : Tensor shape U64 ->
1408 |         (denom : Literal shape Nat) ->
1409 |         {auto 0 isSucc : All IsSucc denom} ->
1410 |         Tensor shape U64
1411 |   rem x y with (x)
1412 |     _ | (MkTensor {shape = _} _) = ewBinary Rem x (tensor {dtype = U64} $ cast <$> y)
1413 |
1414 | export infixr 9 ^
1415 |
1416 | ||| Each element in `base` raised to the power of the corresponding element in `exponent`.
1417 | ||| example, `tensor [2, 25, -9] ^ tensor [3, -0.5, 0.5]` is `tensor [8, 0.2, nan]`.
1418 | |||
1419 | ||| Note: The behaviour of this function is not well-defined at negative or positive infinity, or
1420 | |||   NaN.
1421 | |||
1422 | ||| Note: The first root is used.
1423 | export
1424 | (^) : Tensor shape F64 -> Tensor shape F64 -> Tensor shape F64
1425 | (^) = ewBinary Pow
1426 |
1427 | (>>) : Tensor shape U64 -> Tensor shape U64 -> Tensor shape U64
1428 | (>>) = ewBinary ShiftRightLogical
1429 |
1430 | ||| Element-wise absolute value. For example, `abs (tensor [-2, 3])` is `tensor [2, 3]`.
1431 | export
1432 | abs : Abs dtype => Tensor shape dtype -> Tensor shape dtype
1433 | abs = ewUnary Abs
1434 |
1435 | ||| The element-wise natural exponential. For example, `exp (tensor [-1, 0, 2])` is
1436 | ||| `tensor [1 / euler, 1, pow euler 2]`.
1437 | export
1438 | exp : Tensor shape F64 -> Tensor shape F64
1439 | exp = ewUnary Exp
1440 |
1441 | ||| The element-wise floor function. For example,
1442 | ||| `floor (tensor [-1.6, -1.5, -1.4, -1.0, 1.0, 1.4, 1.5, 1.6])` is
1443 | ||| `tensor [-2.0, -2.0, -2.0, -1.0, 1.0, 1.0, 1.0, 1.0]`.
1444 | export
1445 | floor : Tensor shape F64 -> Tensor shape F64
1446 | floor = ewUnary Floor
1447 |
1448 | ||| The element-wise ceiling function. For example,
1449 | ||| `ceil (tensor [-1.6, -1.5, -1.4, -1.0, 1.0, 1.4, 1.5, 1.6])` is
1450 | ||| `tensor [-1.0, -1.0, -1.0, -1.0, 1.0, 2.0, 2.0, 2.0]`.
1451 | export
1452 | ceil : Tensor shape F64 -> Tensor shape F64
1453 | ceil = ewUnary Ceil
1454 |
1455 | ||| The element-wise natural logarithm. Negative inputs yield NaN output. For example,
1456 | ||| `log (tensor [1 / euler, 1, euler * euler])` is `tensor [-1, 0, 2]`.
1457 | export
1458 | log : Tensor shape F64 -> Tensor shape F64
1459 | log = ewUnary Log
1460 |
1461 | ||| The element-wise logistic function equivalent to `1 / 1 + exp (-x)`.
1462 | export
1463 | logistic : Tensor shape F64 -> Tensor shape F64
1464 | logistic = ewUnary Logistic
1465 |
1466 | ||| The element-wise sine.
1467 | export
1468 | sin : Tensor shape F64 -> Tensor shape F64
1469 | sin = ewUnary Sin
1470 |
1471 | ||| The element-wise cosine.
1472 | export
1473 | cos : Tensor shape F64 -> Tensor shape F64
1474 | cos = ewUnary Cos
1475 |
1476 | ||| The element-wise tangent.
1477 | export
1478 | tan : Tensor shape F64 -> Tensor shape F64
1479 | tan = ewUnary Tan
1480 |
1481 | ||| The element-wise inverse sine.
1482 | export
1483 | asin : Tensor shape F64 -> Tensor shape F64
1484 | asin = ewUnary Asin
1485 |
1486 | ||| The element-wise inverse cosine.
1487 | export
1488 | acos : Tensor shape F64 -> Tensor shape F64
1489 | acos = ewUnary Acos
1490 |
1491 | ||| The element-wise inverse tangent.
1492 | export
1493 | atan : Tensor shape F64 -> Tensor shape F64
1494 | atan = ewUnary Atan
1495 |
1496 | ||| The element-wise hyperbolic sine.
1497 | export
1498 | sinh : Tensor shape F64 -> Tensor shape F64
1499 | sinh = ewUnary Sinh
1500 |
1501 | ||| The element-wise hyperbolic cosine.
1502 | export
1503 | cosh : Tensor shape F64 -> Tensor shape F64
1504 | cosh = ewUnary Cosh
1505 |
1506 | ||| The element-wise hyperbolic tangent.
1507 | export
1508 | tanh : Tensor shape F64 -> Tensor shape F64
1509 | tanh = ewUnary Tanh
1510 |
1511 | ||| The element-wise inverse hyperbolic sine.
1512 | export
1513 | asinh : Tensor shape F64 -> Tensor shape F64
1514 | asinh = ewUnary Asinh
1515 |
1516 | ||| The element-wise inverse hyperbolic cosine.
1517 | export
1518 | acosh : Tensor shape F64 -> Tensor shape F64
1519 | acosh = ewUnary Acosh
1520 |
1521 | ||| The element-wise inverse hyperbolic tangent.
1522 | export
1523 | atanh : Tensor shape F64 -> Tensor shape F64
1524 | atanh = ewUnary Atanh
1525 |
1526 | ||| An approximation to the element-wise error function.
1527 | export
1528 | erf : Tensor shape F64 -> Tensor shape F64
1529 | erf = ewUnary Erf
1530 |
1531 | erfInv : Tensor shape F64 -> Tensor shape F64
1532 | erfInv = ewUnary ErfInv
1533 |
1534 | ||| The element-wise square. For example, `square (tensor [-2, 0, 3])`
1535 | ||| is `tensor [4, 0, 9]`.
1536 | export
1537 | square : Tensor shape F64 -> Tensor shape F64
1538 | square = ewUnary Square
1539 |
1540 | ||| The element-wise square root. The first root is used. Negative inputs yield NaN output.
1541 | ||| For example, `sqrt (tensor [0, 9])` is `tensor [0, 3]`.
1542 | export
1543 | sqrt : Tensor shape F64 -> Tensor shape F64
1544 | sqrt = ewUnary Sqrt
1545 |
1546 | ||| The element-wise minimum of the first argument compared to the second. For example,
1547 | ||| `min (tensor [-3, -1, 3]) (tensor [-1, 0, 1])` is `tensor [-3, -1, 1]`.
1548 | export
1549 | min : Ord dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape dtype
1550 | min (MkTensor x) (MkTensor x') = t0 $ BinaryElementwise Min x x'
1551 |
1552 | namespace Semigroup
1553 |   export
1554 |   [Min] {shape : _} -> Ord dtype => Semigroup (Tensor shape dtype) where
1555 |     (<+>) = min
1556 |
1557 | namespace Monoid
1558 |   export
1559 |   [Min] {shape : _} -> {dtype : _} -> Ord dtype =>
1560 |     Monoid (Tensor shape dtype) using Semigroup.Min where
1561 |       neutral = broadcast max
1562 |
1563 | ||| The element-wise maximum of the first argument compared to the second. For example,
1564 | ||| `max (tensor [-3, -1, 3]) (tensor [-1, 0, 1])` is `tensor [-1, 0, 3]`.
1565 | export
1566 | max : Ord dtype => Tensor shape dtype -> Tensor shape dtype -> Tensor shape dtype
1567 | max (MkTensor x) (MkTensor x') = t0 $ BinaryElementwise Max x x'
1568 |
1569 | namespace Semigroup
1570 |   export
1571 |   [Max] Ord dtype => Semigroup (Tensor shape dtype) where
1572 |     (<+>) = max
1573 |
1574 | namespace Monoid
1575 |   export
1576 |   [Max] {shape : _} -> {dtype : _} -> Ord dtype =>
1577 |     Monoid (Tensor shape dtype) using Semigroup.Max where
1578 |       neutral = broadcast min
1579 |
1580 | ||| The diagonal of a matrix as a vector. For example, for
1581 | ||| ```
1582 | ||| x : Tensor [3, 3] S32
1583 | ||| x = tensor [[0, 1, 2],
1584 | |||             [3, 4, 5],
1585 | |||             [6, 7, 8]]
1586 | ||| ```
1587 | ||| `diag x` is `tensor [0, 4, 8]`.
1588 | export
1589 | diag : Num dtype => Prelude.Num (idrisType dtype) => Tensor [n, n] dtype -> Tensor [n] dtype
1590 | diag {n = 0} x@(MkTensor {shape = [0, 0]} _) = reshape x
1591 | diag {n = S n} x@(MkTensor {shape = [S n, S n]} _) = (x * identity) @@ fill 1
1592 |
1593 | argmxx :
1594 |   Ord dtype =>
1595 |   (Tensor [] dtype -> Tensor [] dtype -> Tensor [] PRED) ->
1596 |   Tensor [] dtype ->
1597 |   Tensor [S n] dtype ->
1598 |   Tag $ Tensor [] U64
1599 | argmxx cmp (MkTensor bound) x@(MkTensor {shape = _} _) = do
1600 |   let MkTensor idxs : Tensor [S n] U64 = iota 0
1601 |       MkTensor x = x
1602 |       MkTensor zero = tensor {dtype = U64} 0
1603 |
1604 |       mon :
1605 |         Tensor [] dtype -> Tensor [] U64 ->
1606 |         Tensor [] dtype -> Tensor [] U64 ->
1607 |         Tag (Tensor [] dtype, Tensor [] U64)
1608 |       mon x y x' y' =
1609 |         let useNext = x == x && (cmp x' x || x' /= x')
1610 |          in pure (select useNext x' x, select useNext y' y)
1611 |
1612 |   MkTagT $ do
1613 |     addr <- reserve
1614 |
1615 |     let MkTagT res = mon
1616 |           (MkTensor $ V 0 $ BoundSet addr)
1617 |           (MkTensor $ V 1 $ BoundSet addr)
1618 |           (MkTensor $ V 2 $ BoundSet addr)
1619 |           (MkTensor $ V 3 $ BoundSet addr)
1620 |         (env, (MkTensor m, MkTensor i)) = runState (emptyFrom !get) res
1621 |         argTys = [TensorType [] dtype, TensorType [] U64]
1622 |         f = MkFn addr (argTys ++ argTys) argTys [m, i] env
1623 |
1624 |     updateCounterFrom env
1625 |     pure $ MkTensor $ V 1 $ Concrete $ Reduce f [bound, zero] [0] [x, idxs]
1626 |
1627 | ||| The first index of the maximum value in a vector. For example,
1628 | ||| `argmax (tensor [-1, 3, -2, -2, 3])` produces `tensor 1`. If the vector contains NaN values,
1629 | ||| `argmax` returns the index of the first NaN.
1630 | export
1631 | argmax : Ord dtype => Tensor [S n] dtype -> Tag $ Tensor [] U64
1632 | argmax x@(MkTensor {dtype} _) = argmxx (>) min x
1633 |
1634 | ||| The first index of the minimum value in a vector. For example,
1635 | ||| `argmin (tensor [-1, 3, -2, -2, 3])` produces `tensor 2`. If the vector contains NaN values,
1636 | ||| `argmin` returns the index of the first NaN.
1637 | export
1638 | argmin : Ord dtype => Tensor [S n] dtype -> Tag $ Tensor [] U64
1639 | argmin x@(MkTensor {dtype} _) = argmxx (<) max x
1640 |
1641 | ||| Represents the upper- or lower-triangular component of a matrix.
1642 | public export
1643 | data Triangle = Upper | Lower
1644 |
1645 | ||| Get the upper- or lower-triangular component of a matrix, always including the matrix diagonal.
1646 | ||| Remaining elements will be zero. For example, for
1647 | ||| ```
1648 | ||| x : Tensor [3, 3] S32
1649 | ||| x = tensor [[1, 2, 3],
1650 | |||             [4, 5, 6],
1651 | |||             [7, 8, 9]]
1652 | ||| ```
1653 | ||| `triangle Lower x` produces
1654 | ||| ```
1655 | ||| x : Tensor [3, 3] S32
1656 | ||| x = tensor [[1, 0, 0],
1657 | |||             [4, 5, 0],
1658 | |||             [7, 8, 9]]
1659 | ||| ```
1660 | export
1661 | triangle :
1662 |   Prelude.Num (idrisType dtype) =>
1663 |   Triangle ->
1664 |   Tensor [n, n] dtype ->
1665 |   Tag $ Tensor [n, n] dtype
1666 | triangle tri (MkTensor x) = do
1667 |   let range : Tensor [n * n] U64 = iota 0
1668 |   indices <- tag $ reshape {to = [n, n], sizesEqual = productSquare n} range
1669 |   let op = case tri of
1670 |         Upper => Tensor.(>)
1671 |         Lower => Tensor.(<)
1672 |   pure $ select (op indices indices.T) (fill $ fromInteger 0) (MkTensor x)
1673 |
1674 |   where
1675 |
1676 |   productSquare : (m : Nat) -> product (the Shape [m * m]) = product (the Shape [m, m])
1677 |   productSquare m =
1678 |     Calc $
1679 |       |~ (m * m) + 0
1680 |       ~~ m * m ... plusZeroRightNeutral (m * m)
1681 |       ~~ (m + 0) * m ..< cong (* m) (plusZeroRightNeutral m)
1682 |
1683 | ||| Cholesky decomposition. Computes the lower triangular matrix `L` from the symmetric, positive
1684 | ||| semi-definite matrix `X` s.t. `X = L @@ L.T`. Values will be NaN if the input matrix is not
1685 | ||| positive semi-definite. The remaining matrix components - those not in the lower triangle or
1686 | ||| diagonal - will always be zero.
1687 | export
1688 | cholesky : Tensor [S n, S n] F64 -> Tag $ Tensor [S n, S n] F64
1689 | cholesky $ MkTensor x = triangle Lower (t0 $ Cholesky x)
1690 |
1691 | export infix 9 |\, \|
1692 |
1693 | namespace Matrix
1694 |   ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is a lower-triangular matrix.
1695 |   ||| `a` is given by the lower-triangular elements of the first argument. Values in the
1696 |   ||| upper-triangular part are ignored. If `a` is lower-triangular already,
1697 |   ||| this is written `a |\ b`.
1698 |   |||
1699 |   ||| The operator is shaped like the lower-triangular portion of a matrix to signal that it uses
1700 |   ||| this portion of its argument. This is in contrast to `(\|)`.
1701 |   export
1702 |   (|\) : Tensor [m, m] F64 -> Tensor [m, n] F64 -> Tensor [m, n] F64
1703 |   (MkTensor a) |\ (MkTensor b) = t0 $ TriangularSolve a b True
1704 |
1705 |   ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is an upper-triangular
1706 |   ||| matrix. `a` is given by the upper-triangular elements of the first argument. Values in the
1707 |   ||| lower-triangular part are ignored. If `a` is upper-triangular already, this is written
1708 |   ||| `a \| b`.
1709 |   |||
1710 |   ||| The operator is shaped like the upper-triangular portion of a matrix to signal that it uses
1711 |   ||| this portion of its argument. This is in contrast to `(|\)`.
1712 |   export
1713 |   (\|) : Tensor [m, m] F64 -> Tensor [m, n] F64 -> Tensor [m, n] F64
1714 |   (MkTensor a) \| (MkTensor b) = t0 $ TriangularSolve a b False
1715 |
1716 | namespace Vector
1717 |   ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is a lower-triangular matrix.
1718 |   ||| `a` is given by the lower-triangular elements of the first argument. Values in the
1719 |   ||| upper-triangular part are ignored. If `a` is lower-triangular already,
1720 |   ||| this is written `a |\ b`.
1721 |   |||
1722 |   ||| The operator is shaped like the lower-triangular portion of a matrix to signal that it uses
1723 |   ||| this portion of its argument. This is in contrast to `(\|)`.
1724 |   export
1725 |   (|\) : Tensor [m, m] F64 -> Tensor [m] F64 -> Tensor [m] F64
1726 |   a |\ b = let (MkTensor {shape = [_]} _) = b in squeeze (a |\ expand 1 b)
1727 |
1728 |   ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is an upper-triangular
1729 |   ||| matrix. `a` is given by the upper-triangular elements of the first argument. Values in the
1730 |   ||| lower-triangular part are ignored. If `a` is upper-triangular already, this is written
1731 |   ||| `a \| b`.
1732 |   |||
1733 |   ||| The operator is shaped like the upper-triangular portion of a matrix to signal that it uses
1734 |   ||| this portion of its argument. This is in contrast to `(|\)`.
1735 |   export
1736 |   (\|) : Tensor [m, m] F64 -> Tensor [m] F64 -> Tensor [m] F64
1737 |   a \| b = let (MkTensor {shape = [_]} _) = b in squeeze (a \| expand 1 b)
1738 |
1739 | ||| Sum the elements along the diagonal of the input. For example,
1740 | ||| `trace (tensor [[-1, 5], [1, 4]])` produces `3`.
1741 | export
1742 | trace : Num dtype => Prelude.Num (idrisType dtype) =>
1743 |         Tensor [S n, S n] dtype ->
1744 |         Tag $ Tensor [] dtype
1745 | trace x with (x)
1746 |   _ | MkTensor {shape = [_, _]} _ = reduce @{Sum} [0, 1] $ x * identity
1747 |
1748 | ||| A `Rand a` produces a pseudo-random value of type `a` from a `Tensor [2] U64` state.
1749 | ||| The state is updated every time a new value is generated.
1750 | public export 0
1751 | Rand : Type -> Type
1752 | Rand = StateT (Tensor [2] U64) Tag
1753 |
1754 | ||| Generate independent and identically distributed (IID) uniform samples.
1755 | |||
1756 | ||| The generated samples are a deterministic function of the input key and state, but may vary
1757 | ||| between PJRT plugin and library version.
1758 | |||
1759 | ||| Example usage, multiplying two uniform samples
1760 | ||| ```
1761 | ||| x : Tag $ Tensor [3] U64
1762 | ||| x = let seed = tensor [1, 1] in evalStateT seed [| rng * rng |]
1763 | ||| ```
1764 | export
1765 | rng : {shape : _} -> Rand $ Tensor shape U64
1766 | rng = ST $ \(MkTensor state) => do
1767 |   res <- tag $ Concrete $ Rng state (TensorType shape U64)
1768 |   pure (MkTensor $ V 0 res, MkTensor $ V 1 res)
1769 |
1770 | ||| Generate independent and identically distributed (IID) from the uniform distribution U(0, 1).
1771 | |||
1772 | ||| The generated samples are a deterministic function of the input key and state, but may vary
1773 | ||| between PJRT plugin and library version.
1774 | |||
1775 | ||| Example usage, multiplying two uniform samples
1776 | ||| ```
1777 | ||| x : Rand $ Tensor [3] F64
1778 | ||| x = [| uniform * uniform |]
1779 | ||| ```
1780 | export
1781 | uniform : {shape : _} -> Rand $ Tensor shape F64
1782 | uniform =
1783 |   let numMantissaBits : Bits64 = 52
1784 |       scale = broadcast $ 2.0 ^ tensor (Scalar $ - cast {to = Double} numMantissaBits)
1785 |       shift = fill $ 64 - numMantissaBits
1786 |    in rng {shape} <&> \x => castDtype (x >> shift) * scale
1787 |
1788 | ||| Generate independent and identically distributed (IID) samples from the standard normal
1789 | ||| distribution N(0, 1).
1790 | |||
1791 | ||| The generated samples are a deterministic function of the input key and state, but may vary
1792 | ||| between PJRT plugin and library version.
1793 | |||
1794 | ||| Example usage, multiplying two normal samples
1795 | ||| ```
1796 | ||| x : Rand $ Tensor [3] F64
1797 | ||| x = [| normal * normal |]
1798 | ||| ```
1799 | export
1800 | normal : {shape : _} -> Rand $ Tensor shape F64
1801 | normal = uniform <&> \x => sqrt (broadcast 2.0) * erfInv (broadcast 2.0 * x - broadcast 1.0)
1802 |