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