0 | module Data.Tensor.Utils
  1 |
  2 | import Data.Nat -- Add import for Cast
  3 | import Data.List
  4 | import System.Random
  5 | import Control.Monad.State
  6 |
  7 | import Data.Tensor.Tensor
  8 | import Data.Container.Additive
  9 | import Data.Container.SubTerm
 10 | import Misc
 11 |
 12 |
 13 | {-------------------------------------------------------------------------------
 14 | {-------------------------------------------------------------------------------
 15 | This file defines common tensor utility functions
 16 | Mirrors those found in numpy/pytorch, and includes:
 17 | * zeros
 18 | * ones
 19 | * range
 20 | * size
 21 | * flatten
 22 | * oneHot
 23 | and the corresponding container variants, when they exist.
 24 |
 25 | Naming needs to be made more consistent
 26 |
 27 | -------------------------------------------------------------------------------}
 28 | -------------------------------------------------------------------------------}
 29 |
 30 | namespace CommonNames
 31 |   public export
 32 |   Scalar : (a : Type) -> Type
 33 |   Scalar a = Tensor [] a
 34 |
 35 |   public export
 36 |   Vector : (c : Axis) -> (a : Type) -> Type
 37 |   Vector c a = Tensor [c] a
 38 |   
 39 |   public export
 40 |   Matrix : (row, col : Axis) -> ConsistentWith row [col] =>
 41 |     (a : Type) -> Type
 42 |   Matrix row col a = Tensor [row, col] a
 43 |
 44 | namespace FillZerosOnes
 45 |   ||| Fills a tensor of a given shape with value `x`
 46 |   ||| Simple an alias for `tensorReplicateAt`
 47 |   public export
 48 |   fillAt : {shape : TensorShape rank} ->
 49 |     (tensorCont shape).Shp -> (x : a) -> Tensor shape a
 50 |   fillAt = tensorReplicateAt
 51 |
 52 |   public export
 53 |   fill : {shape : TensorShape rank} ->
 54 |     AllC TensorMonoid shape =>
 55 |     a -> Tensor shape a
 56 |   fill = tensorReplicate
 57 |
 58 |   public export
 59 |   zerosAt : Num a => {shape : TensorShape rank} ->
 60 |     (tensorCont shape).Shp ->
 61 |     Tensor shape a
 62 |   zerosAt s = fillAt s (fromInteger 0)
 63 |
 64 |   public export
 65 |   zeros : Num a => {shape : TensorShape rank} ->
 66 |     AllC TensorMonoid shape => 
 67 |     Tensor shape a
 68 |   zeros = zerosAt tensorUnitShape
 69 |
 70 |   public export
 71 |   onesAt : Num a => {shape : TensorShape rank} ->
 72 |     (tensorCont shape).Shp ->
 73 |     Tensor shape a
 74 |   onesAt s = fillAt s (fromInteger 1)
 75 |
 76 |   public export
 77 |   ones : Num a => {shape : TensorShape rank} ->
 78 |     AllC TensorMonoid shape => 
 79 |     Tensor shape a
 80 |   ones = onesAt tensorUnitShape
 81 |
 82 |   ||| An identity matrix with True on the diagonal and False elsewhere
 83 |   public export
 84 |   identityBool : {0 c : Axis} -> IsCubical c =>
 85 |     Tensor [c, c] Bool
 86 |   identityBool @{MkIsCubical _ n}
 87 |     = outerWith (==) (positions {sh=()}) (positions {sh=()})
 88 |
 89 |   ||| An identity matrix with ones on the diagonal and zeros elsewhere
 90 |   ||| Analogous to numpy.eye
 91 |   public export
 92 |   identity : {0 c : Axis} -> IsCubical c =>
 93 |     Num a => Tensor [c, c] a
 94 |   identity @{MkIsCubical _ n} = fromBool <$> identityBool
 95 |
 96 | namespace Range
 97 |   {-----
 98 |   This one is interesting, as in the cubical case it's effectively a version of 'tabulate' from Naperian functors.
 99 |   The cubical version is implemented first, and it's possible that a more general version of rangeA can be defined for container based tensors, possibly by tabulating contents of each shape respectively
100 |   -----}
101 |   ||| Separate implementation for the case of one vs two arguments
102 |   ||| This allows the typechecker to more easily match the right implementation at call sites, as with only TwoArg implementation the following doesn't compile:
103 |   ||| a = Tensor [6] Double
104 |   ||| a = arange
105 |   ||| Otoh, we preclude calling this without a type specified
106 |   namespace OneArg
107 |     ||| A range of numbers [0, stop>
108 |     public export
109 |     arange : {0 stop : Axis} -> IsCubical stop =>
110 |       Cast Nat a => Tensor [stop] a
111 |     arange @{MkIsCubical _ n} = cast . finToNat <$> positions {sh=()}
112 |
113 |   namespace TwoArgs
114 |     ||| A range of numbers [start, stop>
115 |     public export
116 |     arangeFromTo : {default (TTInternalName ~~> 0) start : Axis} ->
117 |       {0 stop : Axis} ->
118 |       (cStart : IsCubical start) => (cStop : IsCubical stop) =>
119 |       Cast Nat a => Tensor [stop.name ~~> minus (stop.dim) (start.dim)] a
120 |     arangeFromTo {cStart=(MkIsCubical _ n)} {cStop=(MkIsCubical _ m)}
121 |       = cast . (+n) . finToNat <$> positions {sh=()}
122 |
123 | namespace Flip
124 |   ||| Reverse a tensor along a given axis
125 |   public export
126 |   flip : {shape : TensorShape rank} ->
127 |     (axis : Fin rank) ->
128 |     IsCubical (index axis (toVect shape)) =>
129 |     Tensor shape a -> Tensor shape a
130 |
131 |
132 | namespace Concatenate
133 |   ||| Concatenate two tensors along an existing axis, the first one
134 |   ||| TODO extend to allow concatenation along an arbitrary/named axis
135 |   public export
136 |   concat : {shape : TensorShape rank} -> {l : AxisName} ->
137 |     {x, y : Axis} -> IsCubical x => IsCubical y =>
138 |     ConsistentWith (l ~~> x.dim + y.dim) shape =>
139 |     ConsistentWith x shape =>
140 |     ConsistentWith y shape =>
141 |     Tensor (x :: shape) a ->
142 |     Tensor (y :: shape) a ->
143 |     Tensor ((l ~~> x.dim + y.dim) :: shape) a
144 |   concat @{MkIsCubical _ n} @{MkIsCubical _ m} t t'
145 |     = embedTopExt $ extractTopExt t ++ extractTopExt t'
146 |
147 | namespace Size
148 |   {-----
149 |   Can we measure the size of a tensor of containers?
150 |   Likely need to impose an additional constraint that the set of positions is enumerable
151 |   -----}
152 |   ||| Number of elements in a non-cubical tensor
153 |   public export
154 |   size : {shape : TensorShape rank} ->
155 |     Tensor shape a -> Nat
156 |
157 |   namespace Cubical 
158 |     ||| Number of elements in a cubical tensor
159 |     public export
160 |     size : {shape : TensorShape rank} ->
161 |       All IsCubical (toVect shape) =>
162 |       (0 _ : Tensor shape a) -> Nat
163 |     size {shape} _ = size (toVect shape)
164 |
165 | namespace Flatten
166 |   ||| Flatten a non-cubical tensor into a list
167 |   ||| Requires that we have Foldable on all the components
168 |   ||| In general we won't know the number of elements of a non-cubical tensor at compile time
169 |   public export
170 |   flatten : {0 shape : TensorShape rank} ->
171 |     Foldable (Tensor shape) =>
172 |     Tensor shape a -> List a
173 |   flatten = toList
174 |
175 |   namespace Cubical
176 |     ||| Flatten a cubical tensor into a vector
177 |     ||| Number of elements is known at compile time
178 |     ||| Can even be zero, if any of shape elements is zero
179 |     flatten : {shape : TensorShape rank} ->
180 |       All IsCubical (conts shape) =>
181 |       Tensor shape a -> Vect (size shape) a
182 |     flatten = toVect . extMap (flattenCubical DefaultLayoutOrder) . GetT
183 |
184 | namespace Max
185 |   ||| Maximum value in a tensor
186 |   ||| Returns Nothing if the tensor is empty
187 |   public export
188 |   max : {0 shape : TensorShape rank} ->
189 |     Foldable (Tensor shape) => Ord a =>
190 |     Tensor shape a -> Maybe a
191 |   max = max . flatten
192 |
193 | namespace ArgMinMax
194 |   public export
195 |   argmax : {i : Axis} -> Ord a => IsFoldable i.cont => IsNonEmpty i.cont =>
196 |     (t : Tensor [i] a) -> i.cont.Pos t.extractShapeRank1
197 |   argmax t = argmax (rank1ToExt t)
198 |
199 |   ||| At the moment this simply reuses the vect implementation
200 |   ||| To be revised at some point later
201 |   public export
202 |   argmin : {i : Axis} -> Ord a => IsFoldable i.cont => IsNonEmpty i.cont =>
203 |     (t : Tensor [i] a) -> i.cont.Pos t.extractShapeRank1
204 |   argmin = argmax @{Reverse}
205 |
206 | namespace AllClose
207 |   ||| Scalar approximate equality, following NumPy's `isclose`:
208 |   ||| |x - y| <= atol + rtol * |y|
209 |   public export
210 |   isClose : {default 1.0e-8 atol : Double} -> {default 1.0e-5 rtol : Double} ->
211 |     Double -> Double -> Bool
212 |   isClose x y = abs (x - y) <= atol + rtol * abs y
213 |
214 |   ||| Elementwise approximate equality of tensors, following NumPy's `allclose`
215 |   public export
216 |   allClose : {0 shape : TensorShape rank} ->
217 |     Applicative (Tensor shape) => Foldable (Tensor shape) =>
218 |     {default 1.0e-8 atol : Double} -> {default 1.0e-5 rtol : Double} ->
219 |     Tensor shape Double -> Tensor shape Double -> Bool
220 |   allClose t t' = all id [| isClose {atol} {rtol} t t' |]
221 |
222 | namespace OneHot
223 |   public export
224 |   oneHot : {0 c : Axis} -> IsCubical c =>
225 |     (i : Fin (c.dim)) ->
226 |     Num a =>  Tensor [c] a
227 |   oneHot @{MkIsCubical _ n} i = set zeros [i] 1 
228 |
229 | namespace Triangular
230 |   -- should we have ni,ni here, or ni,nj?
231 |   public export
232 |   cTriBool : {c : Axis} ->
233 |     (ip : InterfaceOnPositions c.cont MOrd) =>
234 |     TensorMonoid c.cont =>
235 |     (sh : c.cont.Shp) -> Tensor [c, c] Bool
236 |   cTriBool sh
237 |     = let cPositions = positions {sh=sh}
238 |           pp : MOrd (c.cont.Pos sh) := GetInterface ip sh
239 |       in outerWith (flip isSubTerm) cPositions cPositions
240 |
241 |   public export
242 |   triBool : {0 c : Axis} -> IsCubical c =>
243 |     Tensor [c, c] Bool
244 |   triBool @{MkIsCubical _ n} = cTriBool ()
245 |
246 |
247 |   ||| A matrix with ones on and below the diagonal, and zeros elsewhere
248 |   ||| Analogous to numpy.tri
249 |   public export
250 |   tri : {0 c : Axis} -> IsCubical c =>
251 |     Num a => Tensor [c, c] a
252 |   tri @{MkIsCubical _ n} = fromBool <$> triBool
253 |
254 |   ||| Lower triangular part of a matrix. Elements above the diagonal are set to
255 |   ||| zero. Analogous to numpy.tril
256 |   public export
257 |   lowerTriangular : {0 c : Axis} -> IsCubical c =>
258 |     Num a => Tensor [c, c] a -> Tensor [c, c] a
259 |   lowerTriangular @{MkIsCubical _ n} = (* tri)
260 |
261 |   ||| Upper triangular part of a matrix. Elements below the diagonal are set to
262 |   ||| zero. Analogous to numpy.triu(.., k=1)
263 |   public export
264 |   upperTriangular : {0 c : Axis} -> IsCubical c =>
265 |     Num a => Tensor [c, c] a -> Tensor [c, c] a
266 |   upperTriangular @{MkIsCubical _ n} = (* ((fromBool . not) <$> triBool))
267 |
268 |   ||| Fill the elements of a tensor `t` with `fill` where `mask` is True
269 |   public export
270 |   maskedFill : {shape : TensorShape rank} ->
271 |     Num a => AllC TensorMonoid shape =>
272 |     (t : Tensor shape a) ->
273 |     (mask : Tensor shape Bool) ->
274 |     (fill : a) ->
275 |     Tensor shape a
276 |   maskedFill t mask fill = liftA2Tensor mask t <&>
277 |     (\(maskVal, tVal) => if maskVal then fill else tVal)
278 |
279 | public export
280 | sum : {shape : TensorShape rank} ->
281 |   Algebra (Tensor shape) a =>
282 |   Tensor shape a -> a
283 | sum = reduce
284 |
285 | public export
286 | mean : {shape : TensorShape rank} ->
287 |   All IsCubical (toVect shape) =>
288 |   Cast Nat a =>
289 |   Fractional a => 
290 |   Algebra (Tensor shape) a =>
291 |   Tensor shape a -> a
292 | mean t = sum t / cast (Cubical.size t)
293 |
294 | public export
295 | variance : {c : Axis} -> IsCubical c =>
296 |   Neg a => Fractional a => Cast Nat a =>
297 |   Tensor [c] a -> a
298 | variance @{MkIsCubical _ n} t =
299 |   let inputMinusMean = t - pure (mean t)
300 |   in mean (inputMinusMean * inputMinusMean)
301 |
302 | -- TODO can we be more pragmatic instead of manually wrapping and unwrapping?
303 | public export
304 | cumulativeSum : {c : Axis} -> Num a =>
305 |   Traversable (Ext c.cont) =>
306 |   Tensor [c] a -> Tensor [c] a
307 | cumulativeSum = extToRank1 . cumulativeSum . rank1ToExt
308 |
309 |
310 |
311 |
312 | namespace TensorComMonoid
313 |   ||| Pointwise commutative monoid structure on tensors, lifted from the one
314 |   ||| on the underlying type
315 |   public export
316 |   tensorComMonoid : {shape : TensorShape rank} -> AllC TensorMonoid shape =>
317 |     ComMonoid a -> ComMonoid (Tensor shape a)
318 |   tensorComMonoid (MkComMonoid p z) = MkComMonoid
319 |     (\t, t' => [| p t t' |])
320 |     (pure z)
321 |
322 | namespace Traversals
323 |   public export
324 |   inorder : Tensor [b ~> BinTreeNode] a -> Tensor [l ~> List] a
325 |   inorder = extToRank1 . extMap BinTreeNode.inorder . rank1ToExt
326 |
327 | namespace Random
328 |   ||| Sampling a tensor is sampling each entry
329 |   ||| Probably can be done more efficiently
330 |   public export
331 |   {shape : TensorShape rank} ->
332 |   Random a =>
333 |   AllC TensorMonoid shape =>
334 |   Traversable (Tensor shape) =>
335 |   Random (Tensor shape a) where
336 |     randomIO = sequence (pure randomIO)
337 |     randomRIO (lo, hi) = traverse randomRIO [| MkPair lo hi |]
338 |
339 |
340 |   tta : Applicative (Tensor ["a" ~~> 1])
341 |   tta = %search
342 |
343 |   ttt : Traversable (Tensor ["b" ~~> 1])
344 |   ttt = %search
345 |   
346 |   ttd : Random Double
347 |   ttd = %search
348 |
349 |   randTensorShape : Random (Tensor ["a" ~~> 2, "b" ~~> 3] Double)
350 |   randTensorShape = %search
351 |
352 | -- Idris can't find the parametric randomIO interface so reimpementing here
353 | public export
354 | random : Num a => Random a => HasIO io =>
355 |   (shape : TensorShape rank) ->
356 |   All IsCubical (toVect shape) =>
357 |   Applicative (Tensor shape) => 
358 |   Traversable (Tensor shape) => 
359 |   io (Tensor shape a)
360 | random shape = sequence $ pure $ randomRIO (0, 1)
361 |