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