0 | {--
1 | Copyright (C) 2021 Joel Berkeley
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.
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.
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`.
42 | ||| A scalar or array. Construct a `Tensor` with function `tensor`.
43 | export
53 | ||| The effect of tagging nodes in a computational graph.
54 | export
62 | export
66 | export
71 | export
75 | export
81 | ||| Mark an expression to be efficiently reused. For example, in
82 | ||| ```
83 | ||| bad : Tensor [9999999] F64
84 | ||| bad = let x = fill {shape = [9999999]} 1.0 in x + x
85 | |||
86 | ||| good : Tag $ Tensor [9999999] F64
87 | ||| good = do x <- tag $ fill {shape = [9999999]} 1.0
88 | ||| pure (x + x)
89 | ||| ```
90 | ||| the large vector `x` is calculated twice in `bad`, but once in `good`, as `tag` marks it for
91 | ||| sharing.
92 | |||
93 | ||| Types that implement this interface should `tag` constituent components it deems worth sharing.
94 | ||| For example, see the implementation for tuples.
95 | |||
96 | ||| See tutorial _Nuisances in the Tensor API_ for details.
102 | export
106 | export
116 | ||| Construct a `Tensor` from `Literal` data. For example
117 | ||| ```
118 | ||| x : Tensor [2, 3] S32
119 | ||| x = tensor [[1, 2, 3],
120 | ||| [4, 5, 6]]
121 | ||| ```
122 | export
127 | export
134 | %hide Literal.All2.All2
138 | ||| Evaluate a list of `Tensor`s as a list of `Literal`s. Tensors in the list can have different
139 | ||| shapes and element types. For example,
140 | ||| ```
141 | ||| main : Device -> IO ()
142 | ||| main device = do [x, y] <- eval device $ do let x = tensor {dtype = F64} [1.2, 3.4]
143 | ||| y <- reduce @{Sum} [0] x
144 | ||| pure [x, y]
145 | ||| printLn x
146 | ||| printLn y
147 | ||| ```
148 | ||| In contrast to `Tensor.eval` when called on multiple tensors, this function constructs and
149 | ||| compiles the graph just once.
162 | where
174 | ||| A convenience wrapper for `List.Tag.eval`, for use with a bare list of `Tensor`s.
184 | ||| Evaluate a `Tensor`, returning its value as a `Literal`. This function builds and executes the
185 | ||| computational graph.
186 | |||
187 | ||| **Note:** Each call to `eval` will rebuild and execute the graph; multiple calls to `eval` on
188 | ||| different tensors, even if they are in the same computation, will be treated independently.
189 | ||| To efficiently evaluate multiple tensors at once, use `List.Tag.eval`.
194 | ||| A convenience wrapper for `Tag.eval`, for use with a bare `Tensor`.
199 | ||| A string representation of a tensor graph.
200 | |||
201 | ||| There are no guarantees whatsoever as to the string structure and contents.
211 | ||| Positive infinity.
212 | export
215 | ||| NaN (not a number).
216 | export
220 | ||| Compares less than or equal to any other value (except NaN).
221 | export
233 | ||| Compares greater than or equal to any other value (except NaN).
234 | export
246 | ||| The most negative possible finite float, approx. -1.8e308
247 | export
251 | ||| The most positive possible finite float, approx. 1.8e308
252 | export
256 | ||| Cast the element type. For example, `castDtype (tensor {dtype = S32} [1, -2])` is
257 | ||| `tensor {dtype = F64} [1.0, -2.0]`.
258 | export
262 | ||| A function type. For example `Func [Nat, String] Bool` is `Nat -> String -> Bool`.
273 | forall arity, rshapes, rdtypes . (shapes : Vect arity Shape) -> (dtypes : Vect arity DType) ->
292 | where
320 | covering
337 | where
354 | ||| Function abstraction in the framework IR.
355 | |||
356 | ||| Operations in spidr are, by default, inlined, which can lead to large IRs. `func` abstracts
357 | ||| (and names) a function in the IR. The resulting function will have the exact same semantics as
358 | ||| the input, though may perform differently, depending on ML compiler behaviour.
359 | |||
360 | ||| Implementation note: MLIR (more specifically the "func" dialect) only supports named functions
361 | ||| at the top (or module) level. To achieve this, `func` lifts its function to the top level,
362 | ||| and automatically converts any variable capture to new function arguments. As such, the IR
363 | ||| may appear different to the Idris code from which it derives.
375 | where
380 | ||| (Experimental) function vectorization.
381 | |||
382 | ||| Lift a function on tensors, so that it applies element-wise across the common leading dimension
383 | ||| of its tensor arguments. For example, for
384 | ||| ```
385 | ||| xs : Tensor [2, 3, 3] S32
386 | ||| xs = tensor [[[ 0, 1, 2],
387 | ||| [ 3, 4, 5],
388 | ||| [ 6, 7, 8]],
389 | ||| [[ 9, 10, 11],
390 | ||| [12, 13, 14],
391 | ||| [15, 16, 17]]]
392 | |||
393 | ||| ys : Tensor [2] S32
394 | ||| ys = Tensor [2, -1]
395 | ||| ```
396 | ||| `do !(vmap (\x, y => pure $ y * diag x)) xs ys` produces `tensor [[0, 8, 16], [-9, -13, -17]]`.
397 | |||
398 | ||| **Warning:** `vmap` is experimental, and only implemented for a subset of the tensor API. You
399 | ||| can see approximately which operations are supported on the Enzyme [tracking issue](https://github.com/EnzymeAD/Enzyme-JAX/issues/152).
413 | where
420 | ||| (Experimental) reverse-mode automatic differentiation.
421 | |||
422 | ||| `grad` can be applied repeatedly to obtain higher derivatives, though we do not yet support
423 | ||| derivatives of vector-valued functions.
424 | |||
425 | ||| For example, for
426 | ||| ```
427 | ||| f : Tensor [2] F64 -> Tag $ Tensor [] F64
428 | ||| f x = do
429 | ||| x <- tag x
430 | ||| let (x0, x1) = (slice [at 0] x, slice [at 1] x)
431 | ||| pure $ x0 / x1
432 | ||| ```
433 | ||| `grad f (tensor [3.0, 2.0])` produces `tensor [0.5, -0.75]`.
434 | |||
435 | ||| **Warning:** `grad` is experimental, and only implemented for a subset of the tensor API. You
436 | ||| can see approximately which operations are supported on the Enzyme [tracking issue](https://github.com/EnzymeAD/Enzyme-JAX/issues/88).
438 | grad : (Tensor shape F64 -> Tag $ Tensor [] F64) -> Tensor shape F64 -> Tag $ Tensor shape F64
441 | %hide Prelude.Interfaces.product
453 | export
458 | ||| Reshape a `Tensor`. For example, `reshape {to = [2, 1]} (tensor [3, 4])` is
459 | ||| `tensor [[3], [4]]`. The output can have a different rank to the input.
460 | export
468 | ||| Add a dimension of length one at the specified `axis`. The new dimension will be at the
469 | ||| specified `axis` in the new `Tensor` (as opposed to the original `Tensor`). For example,
470 | ||| `expand 1 $ tensor [[1, 2], [3, 4], [5, 6]]` is `tensor [[[1, 2]], [[3, 4]], [[5, 6]]]`.
471 | export
480 | ||| A `Squeezable from to` constitutes proof that the shape `from` can be squeezed to the
481 | ||| shape `to`. Squeezing is the process of removing any number of dimensions of length one.
484 | ||| Proof that a shape can be squeezed to itself. For example:
485 | |||
486 | ||| [] to []
487 | ||| [3, 4] to [3, 4]
490 | ||| Proof that any dimensions (including those of length 1) can be preserved in the process of
491 | ||| squeezing. For example:
492 | |||
493 | ||| ...
496 | ||| Proof that any dimensions of length one can be squeezed out. For example:
497 | |||
498 | ||| [1, 3, 1, 1, 4] to [3, 4]
501 | ||| Remove dimensions of length one from a `Tensor` such that it has the desired shape. For example:
502 | |||
503 | ||| ```
504 | ||| x : Tensor [2, 1, 3, 1] S32
505 | ||| x = tensor [[[[4], [5], [6]]],
506 | ||| [[[7], [8], [9]]]]
507 | |||
508 | ||| y : Tensor [2, 1, 3] S32
509 | ||| y = squeeze x
510 | ||| ```
511 | ||| is
512 | ||| ```
513 | ||| y : Tensor [2, 1, 3] S32
514 | ||| y = tensor [[[4, 5, 6]],
515 | ||| [[7, 8, 9]]]
516 | ||| ```
517 | export
525 | ||| A `SliceOrIndex d` is a valid slice or index into a dimension of size `d`. See `slice` for
526 | ||| details.
527 | export
539 | ||| Index at `idx`. See `slice` for details.
545 | ||| Index at the specified index. See `slice` for details.
550 | ||| Slice from `from` (inclusive) to `to` (exclusive). See `slice` for details.
560 | ||| Slice `size` elements starting at the specified scalar `U64` index. See `slice` for details.
565 | ||| Slice across all indices along an axis. See `slice` for details.
570 | ||| A `MultiSlice shape` is a valid multi-dimensional slice into a tensor with shape `shape`.
571 | ||| See `slice` for details.
578 | ||| The shape of a tensor produced by slicing with the specified multi-dimensional slice. See
579 | ||| `Tensor.slice` for details.
588 | ||| Slice or index `Tensor` axes. Each axis can be sliced or indexed, and this can be done with
589 | ||| either static (`Nat`) or dynamic (scalar `U64`) indices.
590 | |||
591 | ||| **Static indices**
592 | |||
593 | ||| Static indices are `Nat`s. For example, for
594 | ||| ```
595 | ||| x : Tensor [5, 6] S32
596 | ||| x = tensor [[ 0, 1, 2, 3, 4, 5],
597 | ||| [ 6, 7, 8, 9, 10, 11],
598 | ||| [12, 13, 14, 15, 16, 17],
599 | ||| [18, 19, 20, 21, 22, 23],
600 | ||| [24, 25, 26, 27, 28, 29]]
601 | ||| ```
602 | ||| we can index as `slice [at 1] x` to get
603 | ||| ```
604 | ||| x : Tensor [6] S32
605 | ||| x = tensor [6, 7, 8, 9, 10, 11]
606 | ||| ```
607 | ||| or we can slice as `slice [2.to 4] x` to get
608 | ||| ```
609 | ||| x : Tensor [2, 6] S32
610 | ||| x = tensor [[12, 13, 14, 15, 16, 17],
611 | ||| [18, 19, 20, 21, 22, 23]]
612 | ||| ```
613 | ||| Note that in `2.to 4`, the 2 is inclusive, and the 4 exclusive, so we return indices 2 and 3.
614 | |||
615 | ||| **Dynamic indices**
616 | |||
617 | ||| Dynamic indices are scalar `U64` values, and the API works slightly differently because we
618 | ||| can't know the value of dynamic indices until the graph is executed. For indexing, with scalar
619 | ||| `U64` index `i` in `slice [at i] x`, `i` is clamped to be a valid index into that dimension.
620 | ||| For example, for `i = tensor 1`, `slice [at i] x` is
621 | ||| ```
622 | ||| x : Tensor [6] S32
623 | ||| x = tensor [6, 7, 8, 9, 10, 11]
624 | ||| ```
625 | ||| as in the static case. However, for `i = tensor 10`, `slice [at i] x` returns the last row
626 | ||| ```
627 | ||| x : Tensor [6] S32
628 | ||| x = tensor [24, 25, 26, 27, 28, 29]
629 | ||| ```
630 | ||| We can also slice by specifying a scalar `U64` start index, and a static size, as
631 | ||| `slice [i.size 2] x` with `i = tensor 2` to get
632 | ||| ```
633 | ||| x : Tensor [2, 6] S32
634 | ||| x = tensor [[12, 13, 14, 15, 16, 17],
635 | ||| [18, 19, 20, 21, 22, 23]]
636 | ||| ```
637 | ||| For a given slice `size`, the dynamic start index is clamped such that we always get `size`
638 | ||| elements along that axis. For example, `slice [i.size 2] x` with `i = tensor 4` is
639 | ||| ```
640 | ||| x : Tensor [2, 6] S32
641 | ||| x = tensor [[18, 19, 20, 21, 22, 23],
642 | ||| [24, 25, 26, 27, 28, 29]]
643 | ||| ```
644 | ||| which starts at index 3 rather than index 4.
645 | |||
646 | ||| **Mixed static, dynamic, slicing and indexing**
647 | |||
648 | ||| Each axis can only be sliced or indexed, and must use only static or dynamic indices. However,
649 | ||| across axes, we can mix these four arbitrarily. For example, with `slice [2.to 4, at 1] x` to
650 | ||| get
651 | ||| ```
652 | ||| x : Tensor [2] S32
653 | ||| x = tensor [13, 19]
654 | ||| ```
655 | ||| or with `i = tensor 2` in `slice [at 1, i.size 2] x` to get
656 | ||| ```
657 | ||| x : Tensor [2] S32
658 | ||| x = tensor [7, 8]
659 | ||| ```
660 | |||
661 | ||| Slices and indices apply to the leading axes of the tensor. For trailing axes omitted from the
662 | ||| multi-dimensional slice, the whole of the axis is returned. If we want to slice or index over
663 | ||| later axes and retain all indices in a leading axis, we can use the convenience function `all`,
664 | ||| as `slice [all, at 3] x` to get
665 | ||| ```
666 | ||| x : Tensor [5] S32
667 | ||| x = tensor [[3], [9], [15], [21], [27]]
668 | ||| ```
669 | ||| This is exactly the same as the more manual `slice [0.to 5, at 3] x` and
670 | ||| `slice [(tensor 0).size 5, at 3] x`.
671 | |||
672 | ||| @at The multi-dimensional slices and indices at which to slice the tensor.
673 | export
676 | let x = val0 $ Slice (mapd start (const 0) at) (mapd stop id at) (replicate (length shape) 1) x
677 | -- we shortcut DynamicSlice to allow autodiff for static slicing
681 | where
721 | ||| The starting indices of a slice with shape `sizes`, into a tensor with shape `bounds`.
722 | export
737 | ||| Scalar.
743 | ||| A `Nat` starting index. This is statically restricted to be in bounds.
753 | ||| A scalar `U64` starting index. This is dynamically truncated to be in bounds.
762 | ||| Replace a slice of a tensor. For example, for
763 | ||| ```
764 | ||| target : Tensor [3, 4] S32
765 | ||| target = tensor [[ 0, 1, 2, 3],
766 | ||| [ 4, 5, 6, 7],
767 | ||| [ 8, 9, 10, 11]]
768 | |||
769 | ||| update : Tensor [2, 2] S32
770 | ||| update = tensor [[12, 13],
771 | ||| [14, 15]]
772 | ||| ```
773 | ||| `updateSlice [0, 1] update target` is
774 | ||| ```
775 | ||| y : Tensor [3, 4] S32
776 | ||| y = tensor [[ 0, 12, 13, 3],
777 | ||| [ 4, 14, 15, 7],
778 | ||| [ 8, 9, 10, 11]]
779 | ||| ```
780 | ||| The starting index can be specified along each axis using either a `Nat` or a `U64` scalar.
781 | ||| Note that the updated slice will always be replaced by `update`, and lie fully within `target`.
782 | ||| This is checked statically for `Nat`. However, for `U64` scalar, the start index `at` is
783 | ||| truncated to the maximum index for which the slice remains within `target`. For example,
784 | ||| `updateSlice [0, (tensor $ Scalar 2)] update target` and
785 | ||| `updateSlice [0, (tensor $ Scalar 3)] update target`
786 | ||| ```
787 | ||| are both
788 | ||| ```
789 | ||| y : Tensor [3, 4] S32
790 | ||| y = tensor [[ 0, 1, 12, 13],
791 | ||| [ 4, 5, 14, 15],
792 | ||| [ 8, 9, 10, 11]]
793 | ||| ```
794 | |||
795 | ||| @at The starting indices of the slice to replace.
796 | ||| @update The tensor to replace the slice with.
797 | ||| @target The tensor in which to replace the slice.
798 | export
807 | where
836 | ||| Concatenate tensors along a common axis. For example,
837 | ||| `concat {leading = []} [tensor [[1, 2]], tensor [[3, 4]], tensor [[5, 6]]]` and
838 | ||| `concat {leading = [_]} [tensor [[1], [3], [5]], tensor [[2], [4], [6]]]` are both
839 | ||| `tensor [[1, 2], [3, 4], [5, 6]]`.
840 | export
845 | -- much of this function is a workaround for https://github.com/idris-lang/Idris2/issues/3860
850 | in t0 {shape} $ Concat (length leading) (x.value :: fromList (forget $ mapProperty (.value) xs))
852 | ||| Transpose a matrix. For example, `(tensor [[1, 2], [3, 4]]).T` is `tensor [[1, 3], [2, 4]]`.
853 | export
857 | ||| Transpose axes of a tensor. This is a more general version of `(.T)`, in which you can
858 | ||| transpose any number of axes in a tensor of arbitrary rank. The i'th axis in the resulting
859 | ||| tensor corresponds to the `index i ordering`'th axis in the input tensor. For example, for
860 | ||| ```
861 | ||| x : Tensor [2, 3, 4] S32
862 | ||| x = tensor [[[ 0, 1, 2, 3],
863 | ||| [ 4, 5, 6, 7],
864 | ||| [ 8, 9, 10, 11]],
865 | ||| [[12, 13, 14, 15],
866 | ||| [16, 17, 18, 19],
867 | ||| [20, 21, 22, 23]]]
868 | ||| ```
869 | ||| `transpose [0, 2, 1] x` is
870 | ||| ```
871 | ||| x : Tensor [2, 4, 3] S32
872 | ||| x = tensor [[[ 0, 4, 8],
873 | ||| [ 1, 5, 9],
874 | ||| [ 2, 6, 10],
875 | ||| [ 3, 7, 11]],
876 | ||| [[12, 16, 20],
877 | ||| [13, 17, 21],
878 | ||| [14, 18, 22],
879 | ||| [15, 19, 23]]]
880 | ||| ```
881 | ||| `transpose [2, 0, 1] x` is
882 | ||| ```
883 | ||| x : Tensor [4, 2, 3] S32
884 | ||| x = tensor [[[ 0, 4, 8],
885 | ||| [12, 16, 20]],
886 | ||| [[ 1, 5, 9],
887 | ||| [13, 17, 21]],
888 | ||| [[ 2, 6, 10],
889 | ||| [14, 18, 22]],
890 | ||| [[ 3, 7, 11],
891 | ||| [15, 19, 23]]]
892 | ||| ```
893 | |||
894 | ||| In order to see what effect transposing a tensor has, it can help to bear in mind the following:
895 | ||| * if an element can be found with `slice [at 3, at 4, at 5] x` in the original tensor,
896 | ||| that same element can instead be found with `slice [at 5, at 3, at 4]` given a
897 | ||| `transpose [2, 0, 1]`. That is, transposing axes re-orders indices when indexing.
898 | ||| * with `transpose [2, 0, 1]`, traversing the first axis in the result is equivalent to
899 | ||| traversing the last axis in the input. Similarly, traversing the last axis in the result is
900 | ||| equivalent to traversing the second axis in the input.
901 | export
911 | ||| A `DimBroadcastable from to` proves that a dimension of size `from` can be broadcast to a
912 | ||| dimension of size `to`.
915 | ||| Proof that any dimension can be broadcast to itself. For example in shapes `[2, 3]` to
916 | ||| `[2, 3]`.
919 | ||| Proof that a dimension of length one can be broadcast to any size. For example in shapes
920 | ||| `[2, 1]` to `[2, 3]`
923 | ||| Proof that any dimension can be broadcast to zero. For example in shapes `[2, 3]` to `[2, 0]`.
927 | ||| A `Broadcastable from to` constitutes proof that the shape `from` can be broadcast to the
928 | ||| shape `to`.
931 | ||| Proof that a shape can be broadcast to itself. For example:
932 | |||
933 | ||| [] to []
934 | ||| [3, 4] to [3, 4]
935 | |||
936 | ||| Implementation note: we could have used `Broadcast [] []`, which would have resulted in more
937 | ||| atomic constructors for `Broadcastable`, but the author guesses that this implementation helps
938 | ||| the type checker avoid applications of `Match`.
941 | ||| Proof that a dimension of size `f` can be broadcast to size `t` if these dimensions
942 | ||| are `DimBroadcastable f t`. For example:
943 | |||
944 | ||| [2, 3] to [2, 3]
945 | ||| [2, 1] to [2, 3]
946 | ||| [2, 1] to [2, 0]
953 | ||| Proof that broadcasting can add outer dimensions i.e. nesting. For example:
954 | |||
955 | ||| [3] to [1, 3]
956 | ||| [3] to [5, 3]
959 | ||| A shape can be extended with any number of leading dimensions.
960 | |||
961 | ||| @leading The leading dimensions.
962 | export
967 | ||| A scalar can be broadcast to any shape.
968 | %hint
969 | export
973 | ||| Broadcast a `Tensor` to a new compatible shape. For example,
974 | ||| ```
975 | ||| x : Tensor [2, 3] S32
976 | ||| x = broadcast (tensor [4, 5, 6])
977 | ||| ```
978 | ||| is
979 | ||| ```
980 | ||| x : Tensor [2, 3] S32
981 | ||| x = tensor [[4, 5, 6], [4, 5, 6]]
982 | ||| ```
983 | export
991 | ||| A `Tensor` where every element has the specified value. For example,
992 | ||| ```
993 | ||| fives : Tensor [2, 3] S32
994 | ||| fives = fill 5
995 | ||| ```
996 | ||| is
997 | ||| ```
998 | ||| fives : Tensor [2, 3] S32
999 | ||| fives = tensor [[5, 5, 5],
1000 | ||| [5, 5, 5]]
1001 | ||| ```
1002 | export
1006 | ||| A constant where values increment from zero along the specified `axis`. For example,
1007 | ||| ```
1008 | ||| x : Tensor [3, 5] S32
1009 | ||| x = iota 1
1010 | ||| ```
1011 | ||| is the same as
1012 | ||| ```
1013 | ||| x : Tensor [3, 5] S32
1014 | ||| x = tensor [[0, 1, 2, 3, 4],
1015 | ||| [0, 1, 2, 3, 4],
1016 | ||| [0, 1, 2, 3, 4]]
1017 | ||| ```
1018 | ||| and
1019 | ||| ```
1020 | ||| x : Tensor [3, 5] S32
1021 | ||| x = iota 0
1022 | ||| ```
1023 | ||| is the same as
1024 | ||| ```
1025 | ||| x : Tensor [3, 5] S32
1026 | ||| x = tensor [[0, 0, 0, 0, 0],
1027 | ||| [1, 1, 1, 1, 1],
1028 | ||| [2, 2, 2, 2, 2]]
1029 | ||| ```
1030 | export
1040 | ||| A while-loop operating on a single tensor.
1041 | |||
1042 | ||| `while1` iteratively checks if the tensor satisfies `condition`, and if it does, updates it with
1043 | ||| `body`.
1044 | |||
1045 | ||| **Note:** The XLA plugins impose heuristic-based rules on variable capture in higher-order
1046 | ||| functions. We do not comprehensively understand these rules, but you *might* experience runtime
1047 | ||| errors if you capture variables in `condition` or `body`, especially variables other than
1048 | ||| constants, or variables that depend on parameters to enclosing (StableHLO) scopes.
1049 | |||
1050 | ||| @condition The guard condition for each iteration.
1051 | ||| @body The update step.
1052 | ||| @initial The initial tensor.
1062 | ||| A while-loop operating on two tensors.
1063 | |||
1064 | ||| `while2` iteratively checks if the tensors together satisfy `condition`, and if so, updates them
1065 | ||| with `body`.
1066 | |||
1067 | ||| **Note:** The XLA plugins impose heuristic-based rules on variable capture in higher-order
1068 | ||| functions. We do not comprehensively understand these rules, but you *might* experience runtime
1069 | ||| errors if you capture variables in `condition` or `body`, especially variables other than
1070 | ||| constants, or variables that depend on parameters to enclosing (StableHLO) scopes.
1071 | |||
1072 | ||| @condition The guard condition for each iteration.
1073 | ||| @body The update step.
1074 | ||| @initial One initial tensor.
1075 | ||| @initial' The other initial tensor.
1083 | res <- tag $ Concrete $ While !(mkFn1 [_, _] [_, _] condition) !(mkFn [_, _] [_, _] body) [i, i']
1086 | ||| Reduce elements along one `axis` of a `Tensor` according to a specified `reducer` `Monoid`.
1087 | ||| For example, if `x = tensor [[0, 1, 2], [3, 4, 5]]`, then reduce @{Sum} 0 x` produces
1088 | ||| `tensor [3, 5, 7]`, and `reduce @{Sum} 1 x` produces `tensor [3, 12]`.
1089 | |||
1090 | ||| **Note:** `Semigroup` doesn't use `Tag`, which limits the functions that can be used in
1091 | ||| `reduce`. However, the most commonly used semigroups don't need `Tag`, including `Sum`,
1092 | ||| `Prod`, `Min` and `Max`, so for ergonomics, we have opted to use `Monoid` as is. We can
1093 | ||| provide an overloaded variant if requested.
1094 | |||
1095 | ||| **Note:** The monoid should not depend on values bound to outer scopes, else runtime errors may
1096 | ||| occur.
1097 | |||
1098 | ||| @reducer How to reduce elements along the given `axis`.
1099 | ||| @axis The axis along which to reduce elements.
1100 | export
1116 | ||| Sort the elements of a `Tensor` along a specified `dimension` according to a scalar-wise
1117 | ||| ordering. For sorting function `f`, elements are sorted such that for consecutive sorted
1118 | ||| elements `a` and `b`, either `f a b` is true, or `f a b` *and* `f b a` are false.
1119 | |||
1120 | ||| **Note:** Sorting is not stable, meaning elements that compare equal according the ordering may
1121 | ||| be sorted in a different order to the order they appear in the input.
1122 | |||
1123 | ||| **Note:** `sort` is limited to use comparison function without `Tag`. However, since the most
1124 | ||| commonly-used functions, including (>), (<), (>=), and (<=), don't use `Tag`, we have opted to
1125 | ||| omit it for ergonomics. We can trivially provide an overloaded variant if requested.
1126 | |||
1127 | ||| **Note:** The comparison function should not depend on values bound to outer scopes, else
1128 | ||| runtime errors may occur.
1129 | |||
1130 | ||| For example, for `x = tensor [[1, 6, 4], [3, 2, 5]]`, `sort (<) 0 x` produces
1131 | ||| `tensor [[1, 2, 4], [3, 6, 5]]`, while `sort (<) 1 x` produces
1132 | ||| `tensor [[1, 4, 6], [2, 3, 5]]`.
1133 | export
1143 | ||| Reverse elements along the specified axes. For example, for
1144 | ||| ```
1145 | ||| x : Tensor [2, 3] S32
1146 | ||| x = tensor [[-2, -1, 0],
1147 | ||| [ 1, 2, 3]]
1148 | ||| ```
1149 | ||| `reverse [0] x` is
1150 | ||| ```
1151 | ||| x : Tensor [2, 3] S32
1152 | ||| x = tensor [[ 1, 2, 3],
1153 | ||| [-2, -1, 0]]
1154 | ||| ```
1155 | ||| `reverse [1] x` is
1156 | ||| ```
1157 | ||| x : Tensor [2, 3] S32
1158 | ||| x = tensor [[ 0, -1, -2],
1159 | ||| [ 3, 2, 1]]
1160 | ||| ```
1161 | ||| and `reverse [0, 1] x` is
1162 | ||| ```
1163 | ||| x : Tensor [2, 3] S32
1164 | ||| x = tensor [[ 3, 2, 1],
1165 | ||| [ 0, -1, -2]]
1166 | ||| ```
1167 | |||
1168 | ||| **Note:** This function requires `axes` is ordered simply so that elements are unique.
1169 | ||| The ordering itself is irrelevant to the implementation, but ensures uniqueness without using
1170 | ||| proofs of contradiction that can be difficult for Idris to construct.
1171 | export
1192 | ||| Element-wise equality. For example, `tensor [1, 2] /= tensor [1, 3]` is
1193 | ||| `tensor [True, False]`.
1194 | export
1198 | ||| Element-wise inequality. For example, `tensor [1, 2] /= tensor [1, 3]` is
1199 | ||| `tensor [False, True]`.
1200 | export
1204 | ||| Element-wise less than. For example, `tensor [1, 2, 3] < tensor [2, 2, 2]` is
1205 | ||| `tensor [True, False, False]`.
1206 | export
1210 | ||| Element-wise greater than. For example, `tensor [1, 2, 3] > tensor [2, 2, 2]` is
1211 | ||| `tensor [False, False, True]`.
1212 | export
1216 | ||| Element-wise less than or equal. For example, `tensor [1, 2, 3] <= tensor [2, 2, 2]`
1217 | ||| is `tensor [True, True, False]`.
1218 | export
1222 | ||| Element-wise greater than or equal. For example,
1223 | ||| `tensor [1, 2, 3] >= tensor [2, 2, 2]` is `tensor [False, True, True]`.
1224 | export
1228 | ||| Element-wise boolean and. For example,
1229 | ||| `tensor [True, True, False, False] && tensor [True, False, True, False]` is
1230 | ||| `tensor [True, False, False, False]`.
1231 | export
1236 | export
1241 | export
1245 | ||| Element-wise boolean or. For example,
1246 | ||| `tensor [True, True, False, False] || tensor [True, False, True, False]` is
1247 | ||| `tensor [True, True, True, False]`.
1248 | export
1253 | export
1258 | export
1262 | ||| Element-wise boolean negation. For example, `not (tensor [True, False])` is
1263 | ||| `tensor [False, True]`.
1264 | export
1268 | ||| Choose elements from two `Tensor`s based on a `Tensor` of predicates. For each element in the
1269 | ||| predicates, the output will use the corresponding element from `onTrue` if the element is
1270 | ||| truthy, else the element from `onFalse`. For example, for
1271 | ||| ```
1272 | ||| preds : Tensor [3] PRED
1273 | ||| preds = tensor [False, True, False]
1274 | |||
1275 | ||| onTrue : Tensor [3] S32
1276 | ||| onTrue = tensor [1, 2, 3]
1277 | |||
1278 | ||| onFalse : Tensor [3] S32
1279 | ||| onFalse = tensor [4, 5, 6]
1280 | ||| ```
1281 | ||| `select preds onTrue onFalse` is `tensor [4, 2, 6]`.
1282 | |||
1283 | ||| @onTrue The elements to choose where the predicate elements are truthy.
1284 | ||| @onFalse The elements to choose where the predicate elements are falsy.
1285 | export
1292 | ||| Use a scalar predicate to evaluate one of two branches. If the predicate is truthy,
1293 | ||| evaluate `onTrue`, else `onFalse`. Each branch is evaluated lazily; only one will be
1294 | ||| evaluated.
1295 | |||
1296 | ||| For example, for
1297 | ||| ```
1298 | ||| f : Tensor [] F64 -> Tag $ Tensor [] F64
1299 | ||| f x = if_ (x < 1.0) (pure $ cos x) (do x <- tag x; x * x)
1300 | ||| ```
1301 | ||| `f 0.0` produces `1.0`, and `f 4.0` produces `8.0`.
1302 | |||
1303 | ||| **Note:** Branches to `if_` are interpreted as constant StableHLO functions. The XLA plugins
1304 | ||| impose heuristic-based rules on variable capture in higher-order functions. We do not
1305 | ||| comprehensively understand these rules, but you *might* experience runtime errors if you capture
1306 | ||| variables in the branches, especially variables other than constants, or variables that depend
1307 | ||| on parameters to enclosing (StableHLO) scopes.
1308 | |||
1309 | ||| @onTrue The branch to evaluate if the predicate is truthy.
1310 | ||| @onFalse The branch to evaluate if the predicate is falsy.
1311 | export
1320 | ||| The identity tensor, with inferred shape and element type. For example,
1321 | ||| ```
1322 | ||| x : Tensor [2, 2] S32
1323 | ||| x = identity
1324 | ||| ```
1325 | ||| is
1326 | ||| ```
1327 | ||| x : Tensor [2, 2] S32
1328 | ||| x = tensor [[1, 0],
1329 | ||| [0, 1]]
1330 | ||| ```
1331 | export
1337 | -- see https://www.python.org/dev/peps/pep-0465/#precedence-and-associativity
1341 | ||| Vector dot product with a tensor of any rank. The vector dot product is with the first axis of
1342 | ||| the right-hand side tensor. For example `tensor [0, 1, 2] @@ tensor [-1, -3, -1]` is
1343 | ||| `-1`.
1344 | export
1349 | ||| Matrix multiplication with a matrix or vector. Contraction is along the last axis of the first
1350 | ||| and the first axis of the last. For example,
1351 | ||| ```
1352 | ||| x : Tensor [2, 3] S32
1353 | ||| x = tensor [[-1, -2, -3],
1354 | ||| [ 0, 1, 2]]
1355 | |||
1356 | ||| y : Tensor [3, 1] S32
1357 | ||| y = tensor [[4, 0, 5]]
1358 | |||
1359 | ||| z : Tensor [2, 1] S32
1360 | ||| z = x @@ y
1361 | ||| ```
1362 | ||| is
1363 | ||| ```
1364 | ||| z : Tensor [2, 1] S32
1365 | ||| z = tensor [-19, 10]
1366 | ||| ```
1367 | export
1373 | (MkTensor x) @@ (MkTensor x') = t0 $ DotGeneral [] [] [1] [0] (TensorType (n :: tl) dtype) x x'
1375 | ||| The output shape of a `dotGeneral` operation.
1383 | Shape
1391 | ||| Matrix multiplication.
1392 | |||
1393 | ||| This is a much more general version of `(@@)`, in which you can specify any number of batch
1394 | ||| and contracting axes. Matrix multiplication is done over each contracting axis.
1395 | ||| The operation is vectorized over batch axes. For each contracting axis on the left-hand
1396 | ||| operand, there is one contracting axis on the right-hand operand. These can be different axes
1397 | ||| in each operand. The same is true for each batch axis.
1398 | |||
1399 | ||| For example, we can vectorize over a typical rank-two matrix multiplication as follows: given
1400 | ||| two inputs tensors
1401 | ||| ```
1402 | ||| let x : Tensor [3, 4, 5, 6] F64
1403 | ||| y : Tensor [3, 4, 6, 7] F64
1404 | ||| ```
1405 | ||| we do
1406 | ||| ```
1407 | ||| let z : Tensor [3, 4, 5, 7] F64 = dotGeneral [0, 1] [0, 1] [3] [2] x y
1408 | ||| ```
1409 | ||| Here, we vectorized over the first two axes `[0, 1]`, and do standard matrix multiplication
1410 | ||| over the remaining axes by specifying the axes 3 and 2 respectively as contracting axes. Notice
1411 | ||| how the batch axes appear once each at the start of the output shape, and the contracting axis
1412 | ||| disappears. Remaining axes appear in order from left to right.
1413 | |||
1414 | ||| Note this API is somewhat of a quickfix to bring general matrix multiplication to the tensor
1415 | ||| API. It is not thoroughly tested. Expect it to change in the future.
1416 | export
1435 | ||| Element-wise addition. For example, `tensor [1, 2] + tensor [3, 4]` is
1436 | ||| `tensor [4, 6]`.
1437 | export
1442 | export
1447 | export
1452 | ||| Element-wise negation. For example, `- tensor [1, -2]` is `tensor [-1, 2]`.
1453 | export
1457 | ||| Element-wise subtraction. For example, `tensor [3, 4] - tensor [4, 2]` is
1458 | ||| `tensor [-1, 2]`.
1459 | export
1463 | ||| Element-wise multiplication. For example, `tensor [2, 3] * tensor [4, 5]` is
1464 | ||| `tensor [8, 15]`.
1465 | export
1470 | export
1475 | export
1480 | ||| Element-wise floating point division. For example, `tensor [2, 3] / tensor [4, 5]` is
1481 | ||| `tensor [0.5, 0.6]`.
1482 | export
1489 | ||| Element-wise division of natural numbers. For example,
1490 | ||| `div (tensor [13, 8]) [3, 4]` is `tensor [4, 2]`.
1491 | |||
1492 | ||| **Note:** Broadcasting a single value into a large `Literal` will be slower than for a `Tensor`.
1493 | ||| For this reason, if you're using a single common denominator, consider using `Scalarwise.div`.
1494 | export
1503 | ||| Overload of `div` for common denominator. The denominator is broadcast to match the numerator.
1504 | export
1514 | ||| Element-wise remainder for natural numbers. For example,
1515 | ||| `rem (tensor [13, 8]) [3, 4]` is `tensor [1, 0]`.
1516 | |||
1517 | ||| **Note:** Broadcasting a single value into a large `Literal` will be slower than for a
1518 | ||| `Tensor`. For this reason, if you're using a single common denominator, consider using
1519 | ||| `Scalarwise.rem`.
1520 | export
1529 | ||| Overload of `rem` for a common denominator. The denominator is broadcast to match the
1530 | ||| numerator.
1531 | export
1540 | ||| Each element in `base` raised to the power of the corresponding element in `exponent`.
1541 | ||| example, `pow (tensor [2, 25, -9]) (tensor [3, -0.5, 0.5])` is `tensor [8, 0.2, nan]`.
1542 | |||
1543 | ||| **Note:** The first root is used. Fractional powers of negative numbers return NaN.
1544 | export
1551 | ||| Element-wise, whether a number is *not* infinity or NaN.
1552 | export
1556 | ||| Element-wise sign. For example, `sign (tensor [-2, 0, 2, nan])` is `tensor [-1, 0, 1, nan]`.
1557 | export
1561 | ||| Element-wise absolute value. For example, `abs (tensor [-2, 3])` is `tensor [2, 3]`.
1562 | export
1566 | ||| The element-wise natural exponential. For example, `exp (tensor [-1, 0, 2])` is
1567 | ||| `tensor [1 / euler, 1, pow euler 2]`.
1568 | export
1572 | ||| The element-wise natural exponential, minus one, equivalent to `(`subtract` 1) . exp`.
1573 | export
1577 | ||| The element-wise floor function. For example,
1578 | ||| `floor (tensor [-1.6, -1.5, -1.4, -1.0, 1.0, 1.4, 1.5, 1.6])` is
1579 | ||| `tensor [-2.0, -2.0, -2.0, -1.0, 1.0, 1.0, 1.0, 1.0]`.
1580 | export
1584 | ||| The element-wise ceiling function. For example,
1585 | ||| `ceil (tensor [-1.6, -1.5, -1.4, -1.0, 1.0, 1.4, 1.5, 1.6])` is
1586 | ||| `tensor [-1.0, -1.0, -1.0, -1.0, 1.0, 2.0, 2.0, 2.0]`.
1587 | export
1594 | ||| Round each element to the nearest integer, breaking ties according to the `BreakTies`.
1595 | |||
1596 | ||| `AwayFromZero` breaks ties away from zero, for example
1597 | ||| `round (tensor [-1.6, -1.5, -1.4, -1.0, -0.5, 0.1, 0.5, 1.0, 1.4, 1.5, 1.6])` is
1598 | ||| `tensor [-2.0, -2.0, -1.0, -1.0, -1.0, 0.0, 1.0, 1.0, 1.0, 2.0, 2.0]`.
1599 | |||
1600 | ||| `ToEven` breaks ties to the nearest even, for example,
1601 | ||| `round (tensor [-1.6, -1.5, -1.4, -1.0, -0.5, 0.5, 1.0, 1.4, 1.5, 1.6])` is
1602 | ||| `tensor [-2.0, -2.0, -1.0, -1.0, 0.0, 0.0, 1.0, 1.0, 2.0, 2.0]`
1603 | export
1608 | ||| The element-wise natural logarithm. Negative inputs yield NaN output. For example,
1609 | ||| `log (tensor [1 / euler, 1, euler * euler])` is `tensor [-1, 0, 2]`.
1610 | export
1614 | ||| The element-wise natural logarithm of the input plus one, equivalent to `log . (+ 1)`.
1615 | ||| Inputs less than (-1) yield NaN.
1616 | export
1620 | ||| The element-wise logistic function equivalent to `1 / 1 + exp (-x)`.
1621 | export
1625 | ||| The element-wise sine.
1626 | export
1630 | ||| The element-wise cosine.
1631 | export
1635 | ||| The element-wise tangent.
1636 | export
1640 | ||| The element-wise arcsine.
1641 | export
1645 | ||| The element-wise arccosine.
1646 | export
1650 | ||| The element-wise arctangent.
1651 | export
1655 | ||| The element-wise two-argument arctangent: the angle between the vector (x, y) and the positive
1656 | ||| x-axis.
1657 | export
1661 | ||| The element-wise hyperbolic sine.
1662 | export
1666 | ||| The element-wise hyperbolic cosine.
1667 | export
1671 | ||| The element-wise hyperbolic tangent.
1672 | export
1676 | ||| The element-wise hyperbolic arcsine.
1677 | export
1681 | ||| The element-wise hyperbolic arccosine.
1682 | export
1686 | ||| The element-wise hyperbolic arctangent.
1687 | export
1691 | ||| An approximation to the element-wise error function.
1692 | export
1699 | ||| The element-wise square. For example, `square (tensor [-2, 0, 3])`
1700 | ||| is `tensor [4, 0, 9]`.
1701 | export
1705 | ||| The element-wise square root. The first root is used. Negative inputs yield NaN output.
1706 | ||| For example, `sqrt (tensor [0, 9])` is `tensor [0, 3]`.
1707 | export
1711 | ||| The element-wise reciprocal square root. The first root is used. Negative inputs yield NaN output.
1712 | ||| For example, `rsqrt (tensor [0, 4])` is `tensor [inf, 0.5]`.
1713 | export
1717 | ||| The element-wise cube root. The first root is used. Negative inputs yield negative outputs.
1718 | ||| For example, `cbrt (tensor [8, -27])` is `tensor [2, -3]`.
1719 | export
1723 | ||| The element-wise minimum of the first argument compared to the second. For example,
1724 | ||| `min (tensor [-3, -1, 3]) (tensor [-1, 0, 1])` is `tensor [-3, -1, 1]`.
1725 | export
1730 | export
1735 | export
1740 | ||| The element-wise maximum of the first argument compared to the second. For example,
1741 | ||| `max (tensor [-3, -1, 3]) (tensor [-1, 0, 1])` is `tensor [-1, 0, 3]`.
1742 | export
1747 | export
1752 | export
1757 | ||| The diagonal of a matrix as a vector. For example, for
1758 | ||| ```
1759 | ||| x : Tensor [3, 3] S32
1760 | ||| x = tensor [[0, 1, 2],
1761 | ||| [3, 4, 5],
1762 | ||| [6, 7, 8]]
1763 | ||| ```
1764 | ||| `diag x` is `tensor [0, 4, 8]`.
1765 | export
1804 | ||| The first index of the maximum value in a vector. For example,
1805 | ||| `argmax (tensor [-1, 3, -2, -2, 3])` produces `tensor 1`. If the vector contains NaN values,
1806 | ||| `argmax` returns the index of the first NaN.
1807 | export
1811 | ||| The first index of the minimum value in a vector. For example,
1812 | ||| `argmin (tensor [-1, 3, -2, -2, 3])` produces `tensor 2`. If the vector contains NaN values,
1813 | ||| `argmin` returns the index of the first NaN.
1814 | export
1818 | ||| Represents the upper- or lower-triangular component of a matrix.
1822 | ||| Get the upper- or lower-triangular component of a matrix, always including the matrix diagonal.
1823 | ||| Remaining elements will be zero. For example, for
1824 | ||| ```
1825 | ||| x : Tensor [3, 3] S32
1826 | ||| x = tensor [[1, 2, 3],
1827 | ||| [4, 5, 6],
1828 | ||| [7, 8, 9]]
1829 | ||| ```
1830 | ||| `triangle Lower x` produces
1831 | ||| ```
1832 | ||| x : Tensor [3, 3] S32
1833 | ||| x = tensor [[1, 0, 0],
1834 | ||| [4, 5, 0],
1835 | ||| [7, 8, 9]]
1836 | ||| ```
1837 | export
1851 | where
1857 | Refl
1859 | ||| Cholesky decomposition. Computes the lower triangular matrix `L` from the symmetric, positive
1860 | ||| semi-definite matrix `X` s.t. `X = L @@ L.T`. Values will be NaN if the input matrix is not
1861 | ||| positive semi-definite. The remaining matrix components - those not in the lower triangle or
1862 | ||| diagonal - will always be zero.
1863 | export
1870 | ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is a lower-triangular matrix.
1871 | ||| `a` is given by the lower-triangular elements of the first argument. Values in the
1872 | ||| upper-triangular part are ignored. If `a` is lower-triangular already,
1873 | ||| this is written `a |\ b`.
1874 | |||
1875 | ||| The operator is shaped like the lower-triangular portion of a matrix to signal that it uses
1876 | ||| this portion of its argument. This is in contrast to `(\|)`.
1877 | export
1881 | ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is an upper-triangular
1882 | ||| matrix. `a` is given by the upper-triangular elements of the first argument. Values in the
1883 | ||| lower-triangular part are ignored. If `a` is upper-triangular already, this is written
1884 | ||| `a \| b`.
1885 | |||
1886 | ||| The operator is shaped like the upper-triangular portion of a matrix to signal that it uses
1887 | ||| this portion of its argument. This is in contrast to `(|\)`.
1888 | export
1893 | ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is a lower-triangular matrix.
1894 | ||| `a` is given by the lower-triangular elements of the first argument. Values in the
1895 | ||| upper-triangular part are ignored. If `a` is lower-triangular already,
1896 | ||| this is written `a |\ b`.
1897 | |||
1898 | ||| The operator is shaped like the lower-triangular portion of a matrix to signal that it uses
1899 | ||| this portion of its argument. This is in contrast to `(\|)`.
1900 | export
1904 | ||| Solve the set of linear equations `a @@ x = b` for `x` where `a` is an upper-triangular
1905 | ||| matrix. `a` is given by the upper-triangular elements of the first argument. Values in the
1906 | ||| lower-triangular part are ignored. If `a` is upper-triangular already, this is written
1907 | ||| `a \| b`.
1908 | |||
1909 | ||| The operator is shaped like the upper-triangular portion of a matrix to signal that it uses
1910 | ||| this portion of its argument. This is in contrast to `(|\)`.
1911 | export
1915 | ||| Sum the elements along the diagonal of the input. For example,
1916 | ||| `trace (tensor [[-1, 5], [1, 4]])` produces `3`.
1917 | export
1924 | ||| A `Rand a` produces a pseudo-random value of type `a` from a `Tensor [2] U64` state.
1925 | ||| The state is updated every time a new value is generated.
1930 | ||| Generate independent and identically distributed (IID) uniform samples.
1931 | |||
1932 | ||| The generated samples are a deterministic function of the input key and state, but may vary
1933 | ||| between PJRT plugin and library version.
1934 | |||
1935 | ||| Example usage, multiplying two uniform samples
1936 | ||| ```
1937 | ||| x : Tag $ Tensor [3] U64
1938 | ||| x = let seed = tensor [1, 1] in evalStateT seed [| rng * rng |]
1939 | ||| ```
1940 | export
1946 | ||| Generate independent and identically distributed (IID) from the uniform distribution U(0, 1).
1947 | |||
1948 | ||| The generated samples are a deterministic function of the input key and state, but may vary
1949 | ||| between PJRT plugin and library version.
1950 | |||
1951 | ||| Example usage, multiplying two uniform samples
1952 | ||| ```
1953 | ||| x : Rand $ Tensor [3] F64
1954 | ||| x = [| uniform * uniform |]
1955 | ||| ```
1956 | export
1964 | ||| Generate independent and identically distributed (IID) samples from the standard normal
1965 | ||| distribution N(0, 1).
1966 | |||
1967 | ||| The generated samples are a deterministic function of the input key and state, but may vary
1968 | ||| between PJRT plugin and library version.
1969 | |||
1970 | ||| Example usage, multiplying two normal samples
1971 | ||| ```
1972 | ||| x : Rand $ Tensor [3] F64
1973 | ||| x = [| normal * normal |]
1974 | ||| ```
1975 | export