0 | module Data.Tensor.Shape.Shape
  1 |
  2 | import public Decidable.Equality
  3 | import Data.Vect.Elem
  4 | import Data.Vect.Quantifiers
  5 | import Data.List.Quantifiers
  6 |
  7 | import Data.Container.Base
  8 | import Data.Tensor.Shape.Axis
  9 | import Misc
 10 |
 11 | {-------------------------------------------------------------------------------
 12 | {-------------------------------------------------------------------------------
 13 |
 14 | ~~~~~~~~~~~~~~~
 15 | Design choices:
 16 | ~~~~~~~~~~~~~~~
 17 |
 18 | 1) Persistent axis names.
 19 |
 20 | Instead of transient axis names - where the names are bound within a function such as `np.einsum("ij->ji", m)` and erased after the said function is evaluated - axis names are a part of the tensor shape, and persist with the lifetime of the said tensor.
 21 |
 22 | 2) Axis declarations persist globally, but are only checked for consistency at call sites.
 23 |
 24 | In a proper tensor programming language we'd prevent declaration of inconsistent/duplicate axis names to begin with.
 25 | Here we opt for a more pragmatic approach: checking consistency locally at each call site, rather than at declaration sites. 
 26 |
 27 | 3) Duplicate axis names within a tensor are allowed, as long as the names are consistent.
 28 |
 29 | Otherwise it would not be clear how to take the diagonal/trace of a matrix while referring only to the axes: they'd have to have different names.
 30 |
 31 | 4) Tensor contractions allow duplicate names both in the input and in the output, again, as long as they're consistent. Names in output which haven't appeared in the input are also allowed.
 32 |
 33 | Duplicate input means zipping/reduction, duplicate output means diagonalisation.
 34 | This is different from what standard `einsum` allows: it does not permit duplicate names in the output. It also does not not allow names to appear in output which haven't appeared in the input, because their size wouldn't be known. Becuase for us `Axis` refers to both name and size, this works.
 35 |
 36 | ~~~~~~~~~~~~~~~
 37 | Design choices:
 38 | ~~~~~~~~~~~~~~~
 39 |
 40 | Two axes are consistent if they either have different names, or same name and same underlying container. Conversely, they're inconsistent if they're called the same, but refer to the different underlying container.
 41 |
 42 | It is possible to have inconsistent axes bound declared within the same scope.
 43 | Consistency is checked only at call sites.
 44 |
 45 | Alternatively if we were building a programming languge we'd check consistency with each declaration. That is, writing something like:
 46 | ```idris
 47 | BatchSize1 : Axis
 48 | BatchSize1 = "batchSize" ~> Vect 128
 49 |
 50 | BatchSize2 : Axis
 51 | BatchSize2 = "batchSize" ~> Vect 129
 52 | ```
 53 | would throw an error on the line `BatchSize2 = ...` because we're redeclaring "batchSize" which already exists.
 54 |
 55 | ------------------------------------------- 
 56 |
 57 | Similar projects/ideas:
 58 | * XArray: https://docs.xarray.dev/en/stable/ (persistent axis names)
 59 | * Haliax: https://github.com/marin-community/haliax
 60 |
 61 | Useful documentation:
 62 | * https://obilaniu6266h16.wordpress.com/2016/02/04/einstein-summation-in-numpy/
 63 | * https://nlp.seas.harvard.edu/NamedTensor
 64 | * https://einsum.joelburget.com/
 65 |
 66 |
 67 | -------------------------------------------------------------------------------}
 68 | -------------------------------------------------------------------------------}
 69 |
 70 |
 71 | mutual
 72 |   ||| Datatype defining the shape of a tensor
 73 |   ||| It is a vector of consistently named axes
 74 |   public export
 75 |   data TensorShape : (rank : Nat) ->Type where
 76 |     Nil : TensorShape 0
 77 |     (::) : (a : Axis) -> (as : TensorShape k) ->
 78 |       (axisConsistent : a `ConsistentWith` as) =>
 79 |       TensorShape (S k)
 80 |
 81 |   ||| An axis is named consistently with respect to an existing tensor shape
 82 |   ||| if its name does not appear in the shape, or it if it appears and
 83 |   ||| references the same container
 84 |   public export
 85 |   data ConsistentWith : Axis -> TensorShape k -> Type where
 86 |     NewAxis : {0 a : Axis} -> {0 as : TensorShape k} ->
 87 |       NotElem a.name as ->
 88 |       a `ConsistentWith` as
 89 |     ExistingAxis : {0 a : Axis} -> {0 as : TensorShape k} ->
 90 |       (e : Elem a.name as) ->
 91 |       index as a.name = a.cont ->
 92 |       a `ConsistentWith` as
 93 |  
 94 |   ||| A proof that an axis name `a` is found in a tensor shape `as`
 95 |   ||| Notably this could have been implemented to check whether the 
 96 |   ||| *entire axis* is in vector, and not just the name.
 97 |   ||| But since a tensor shape is more akin to a dictionary, keeping this form
 98 |   public export
 99 |   data Elem : (axisName : AxisName) -> (as : TensorShape rank) -> Type where
100 |     Here : {0 as : TensorShape rank} ->
101 |       a `ConsistentWith` as =>
102 |       axisName = a.name =>
103 |       Elem axisName (a :: as)
104 |     There : {0 a : Axis} -> {0 as : TensorShape rank} ->
105 |       a `ConsistentWith` as =>
106 |       (elem : Elem axisName as) =>
107 |       Elem axisName (a :: as)
108 |
109 |   ||| A proof that an axis name `a` is not found in a tensor shape `as`
110 |   public export
111 |   data NotElem : (axisElem : AxisName) -> (as : TensorShape rank) -> Type where
112 |     NotInEmpty : NotElem axisElem []
113 |     NotInNonEmpty : {0 axisElem : AxisName} -> {0 a : Axis} ->
114 |       {0 as : TensorShape rank} ->
115 |       (neq : IsNo (decEq axisElem a.name)) ->
116 |       (notElem : NotElem axisElem as) =>
117 |       (a `ConsistentWith` as) ->
118 |       NotElem axisElem (a :: as)
119 |
120 |       
121 |   ||| Indexing into a tensor shape.
122 |   ||| Could be many occurences - recovers the one provided by `isElem`
123 |   public export
124 |   index : (shape : TensorShape rank) ->
125 |     (axisName : AxisName) -> 
126 |     (isElem : Elem axisName shape) =>
127 |     Cont
128 |   index (a :: _) axisName @{Here} = a.cont
129 |   index (_ :: as) axisName @{There} = index as axisName
130 |
131 | ||| to get rid of believe_me this might need to be put in a mutual block too
132 | public export
133 | rename : (shape : TensorShape rank) ->
134 |   (axisName : AxisName) ->
135 |   (newAxisName : AxisName) ->
136 |   TensorShape rank
137 | rename [] _ _ = []
138 | rename (a :: as) axisName newAxisName
139 |   = (::) (applyWhen (axisName == a.name) (flip rename newAxisName) a) (rename as axisName newAxisName) @{believe_me "consistentAfterRenaming_believe_me"}
140 |
141 | namespace RenameByIndex
142 |   ||| to get rid of believe_me this might need to be put in a mutual block too
143 |   public export
144 |   rename : (shape : TensorShape rank) ->
145 |     (axisIndex : Fin rank) ->
146 |     (newAxisName : AxisName) ->
147 |     TensorShape rank
148 |   rename (a :: as) FZ newAxisName
149 |     = (::) (newAxisName ~> a.cont) as @{believe_me "consistentAfterRenamingByIndex"}
150 |   rename (a :: as) (FS axisIndex) newAxisName
151 |     = (::) a (rename as axisIndex newAxisName) @{believe_me "consistentAfterRenamingByIndex"}
152 |
153 | ||| These are quantifiers for the axes
154 | ||| Sometimes we need to explicitly quantify over something involving a name
155 | namespace Quantifiers
156 |   public export
157 |   data All : (p : Axis -> Type) -> TensorShape k -> Type where
158 |     Nil : All p []
159 |     (::) : {0 a : Axis} -> {0 as : TensorShape k} ->
160 |       p a -> All p as ->
161 |       a `ConsistentWith` as =>
162 |       All p (a :: as)
163 |
164 |   public export
165 |   data Any : (p : Axis -> Type) -> TensorShape k -> Type where
166 |     Here : {0 a : Axis} -> {0 as : TensorShape k} ->
167 |       a `ConsistentWith` as =>
168 |       p a -> Any p (a :: as)
169 |     There : {0 a : Axis} -> {0 as : TensorShape k} ->
170 |       a `ConsistentWith` as =>
171 |       Any p as -> Any p (a :: as)
172 |   
173 |   ||| These quantifiers are specifically for underlying containers, not axes
174 |   ||| In theory we should be able to overload, but what do we do with 
175 |   ||| purposefully overloaded instances `IsCubical`, `IsNaperian` whose
176 |   ||| existence will prevent elaboration?
177 |   namespace QuantifierOnContainers
178 |     public export
179 |     data AllC : (p : Cont -> Type) -> TensorShape k -> Type where
180 |       Nil : AllC p []
181 |       (::) : {0 a : Axis} -> {0 as : TensorShape k} ->
182 |         p a.cont -> AllC p as ->
183 |         a `ConsistentWith` as =>
184 |         AllC p (a :: as)
185 |
186 |     public export
187 |     data AnyC : (p : Cont -> Type) -> TensorShape k -> Type where
188 |       Here : {0 a : Axis} -> {0 as : TensorShape k} ->
189 |         a `ConsistentWith` as =>
190 |         p a.cont -> AnyC p (a :: as)
191 |       There : {0 a : Axis} -> {0 as : TensorShape k} ->
192 |         a `ConsistentWith` as =>
193 |         AnyC p as -> AnyC p (a :: as)
194 |
195 | public export
196 | tensorShapesConsistent : TensorShape k -> TensorShape k' -> Type
197 | tensorShapesConsistent s1 s2 = All (\a => a `ConsistentWith` s2) s1
198 |
199 | ||| (::) here pattern matches on the proof `axisConsistent` and discards it
200 | public export
201 | toVect : TensorShape k -> Vect k Axis
202 | toVect [] = []
203 | toVect (a :: as) = a :: toVect as
204 |
205 | public export
206 | toList : TensorShape k -> List Axis
207 | toList [] = []
208 | toList (a :: as) = a :: toList as
209 |
210 | ||| Convenience function, turns it also into a list
211 | ||| Because `Tensor` from `Data.Container` uses lists with tensors
212 | public export
213 | conts : TensorShape k -> List Cont
214 | conts ts = cont <$> toList ts
215 |
216 | ||| A property of every axis's container, as a property of the list `conts`
217 | public export
218 | toAll : {0 shape : TensorShape k} ->
219 |   AllC p shape -> List.Quantifiers.All.All p (conts shape)
220 | toAll [] = []
221 | toAll (x :: xs) = x :: toAll xs
222 |
223 | ||| Renaming the shape preserves the underlying data
224 | public export
225 | renamePreservesConts : (shape : TensorShape rank) ->
226 |   (axisName : AxisName) ->
227 |   (newAxisName : AxisName) ->
228 |   conts (rename shape axisName newAxisName) = conts shape
229 | renamePreservesConts [] _ _ = Refl
230 | renamePreservesConts (a :: as) axisName newAxisName with (axisName == a.name)
231 |   _ | True = cong (a.cont ::) (renamePreservesConts as axisName newAxisName)
232 |   _ | False = cong (a.cont ::) (renamePreservesConts as axisName newAxisName)
233 |
234 | namespace RenameByIndex
235 |   ||| Renaming a shape at a specific index preserves the underlying containers.
236 |   public export
237 |   renamePreservesConts : (shape : TensorShape rank) ->
238 |     (axisIndex : Fin rank) ->
239 |     (newAxisName : AxisName) ->
240 |     conts (rename shape axisIndex newAxisName) = conts shape
241 |   renamePreservesConts (a :: as) FZ newAxisName = Refl
242 |   renamePreservesConts (a :: as) (FS axisIndex) newAxisName
243 |     = cong (a.cont ::) (renamePreservesConts as axisIndex newAxisName)
244 |
245 | ||| Names of the axes in a tensor shape
246 | public export
247 | axisNames : TensorShape k -> Vect k AxisName
248 | axisNames ts = name <$> toVect ts
249 |
250 | ||| Sizes of the axes in a tensor shape
251 | public export
252 | axisSizes : TensorShape k -> Vect k Cont
253 | axisSizes ts = cont <$> toVect ts
254 |
255 | ||| Size of a tensor shape, i.e. its number of elements
256 | public export
257 | size : (shape : TensorShape k) -> All IsCubical (conts shape) => Nat
258 | size shape = size (conts shape)
259 |
260 | ||| Cubicality evidence for a tensor shape, using `Either` so that
261 | ||| auto-search tries `Left prf` (positive evidence) before `Right ()`
262 | ||| (fallback). Standard `Maybe` can't be used because `Nothing` is
263 | ||| listed first in its definition and auto-search would always pick it.
264 | public export
265 | 0 TensorCubEvidence : TensorShape k -> Type
266 | TensorCubEvidence shape = Either (All IsCubical shape) ()
267 |
268 |
269 |
270 | namespace Unique
271 |   ||| A proof that an axis name only appears once in the tensor shape
272 |   public export
273 |   data UniqueElem : AxisName -> TensorShape rank -> Type where
274 |     Here : {0 as : TensorShape rank} ->
275 |       axisName = ax.name =>
276 |       NotElem axisName as =>
277 |       ax `ConsistentWith` as =>
278 |       UniqueElem axisName (ax :: as)
279 |     There : {0 ax : Axis} -> {0 as : TensorShape rank} ->
280 |       IsSucc rank =>
281 |       (uniqueElem : UniqueElem axisName as) =>
282 |       (neq : IsNo (decEq axisName ax.name)) =>
283 |       ax `ConsistentWith` as =>
284 |       UniqueElem axisName (ax :: as)
285 |
286 |   ||| Forgets that the axis name only appears once in the tensor shape
287 |   public export
288 |   forgetUnique : {as : TensorShape rank} ->
289 |     UniqueElem axisName as ->
290 |     Elem axisName as
291 |   forgetUnique {as = (a :: as)} Here = Here 
292 |   forgetUnique {as = (a :: as)} (There {uniqueElem=elem})
293 |     = There {elem=forgetUnique elem}
294 |
295 |   public export
296 |   index : (shape : TensorShape rank) ->
297 |     (axisName : AxisName) ->
298 |     (uniqueElem : UniqueElem axisName shape) =>
299 |     Cont
300 |   index (a :: _) axisName @{Here} = a.cont
301 |   index (_ :: as) axisName @{There} = index as axisName
302 |
303 |   mutual
304 |     public export
305 |     removeAxis : {rank : Nat} ->
306 |       (toRemove : AxisName) ->
307 |       (shape : TensorShape (S rank)) ->
308 |       (is : UniqueElem toRemove shape) =>
309 |       TensorShape rank
310 |     removeAxis toRemove (_ :: as) @{Here} = as
311 |     removeAxis toRemove (a :: as) @{There @{ItIsSucc}}
312 |       = let cProof = consistentAfterRemoving a as toRemove
313 |         in a :: removeAxis toRemove as
314 |
315 |     ||| After removing the (unique) axis name `toRemove`, said name no longer
316 |     ||| appears in the shape
317 |     public export
318 |     removedNotElem : {rank : Nat} ->
319 |       (toRemove : AxisName) -> (as : TensorShape (S rank)) ->
320 |       (uElem : UniqueElem toRemove as) =>
321 |       NotElem toRemove (removeAxis toRemove as)
322 |     removedNotElem _ (_ :: _) @{Here @{_} @{tNotElem}} = tNotElem
323 |     removedNotElem toRemove (_ :: as) @{There @{ItIsSucc} @{_} @{neqRm}}
324 |       = NotInNonEmpty neqRm @{removedNotElem toRemove as} _
325 |
326 |     ||| A name absent from a shape is still absent after removing an axis
327 |     public export
328 |     notElemAfterRemoving : {rank : Nat} ->
329 |       (x : AxisName) -> (as : TensorShape (S rank)) ->
330 |       NotElem x as ->
331 |       (toRemove : AxisName) ->
332 |       (uElem : UniqueElem toRemove as) =>
333 |       NotElem x (removeAxis toRemove as)
334 |     notElemAfterRemoving _ (_ :: _) (NotInNonEmpty _ @{notElemTail} _) _ @{Here}
335 |       = notElemTail
336 |     notElemAfterRemoving x (_ :: as) (NotInNonEmpty neqX @{notElemTail} _) toRemove @{There @{ItIsSucc}}
337 |       = NotInNonEmpty neqX @{notElemAfterRemoving x as notElemTail toRemove} _
338 |
339 |     ||| A name different from `toRemove` that appears in a shape still appears
340 |     ||| after the removal, and indexes the same container
341 |     public export
342 |     indexAfterRemoving : {rank : Nat} ->
343 |       (x : AxisName) -> (as : TensorShape (S rank)) ->
344 |       (e : Elem x as) ->
345 |       (toRemove : AxisName) ->
346 |       (uElem : UniqueElem toRemove as) =>
347 |       (neq : IsNo (decEq x toRemove)) =>
348 |       (e' : Elem x (removeAxis toRemove as)
349 |         ** index (removeAxis toRemove as) x = index as x)
350 |     indexAfterRemoving _ ((_ ~> _) :: _) (Here @{_} @{Refl}) _ @{Here @{Refl}} @{neq}
351 |       = absurd @{UninhabitedIsNoRefl} neq
352 |     indexAfterRemoving _ (_ :: _) (There @{_} @{elemTail}) _ @{Here}
353 |       = (elemTail ** Refl)
354 |     indexAfterRemoving _ (_ :: _) (Here @{_} @{eqPrf}) _ @{There @{ItIsSucc}}
355 |       = (Here @{_} @{eqPrf} ** Refl)
356 |     indexAfterRemoving x (ax :: as) (There @{_} @{elemTail}) toRemove @{There @{ItIsSucc}}
357 |       = let (e' ** prfIdx) = indexAfterRemoving x as elemTail toRemove
358 |         in (There @{consistentAfterRemoving ax as toRemove} {elem = e'} ** prfIdx)
359 |
360 |     ||| Proof that an axis `a` is consistent with `as` after removing the axis 
361 |     ||| name `toRemove` from `as`
362 |     public export
363 |     consistentAfterRemoving : {rank : Nat} ->
364 |       (a : Axis) -> (as : TensorShape (S rank)) ->
365 |       a `ConsistentWith` as =>
366 |       (toRemove : AxisName) ->
367 |       (uElem : UniqueElem toRemove as) =>
368 |       a `ConsistentWith` (removeAxis toRemove as)
369 |     consistentAfterRemoving a as @{cw} toRemove with (decEq a.name toRemove)
370 |       _ | Yes eqPrf
371 |           = NewAxis (rewrite eqPrf in removedNotElem toRemove as)
372 |       _ | No contra = case cw of
373 |             NewAxis notElem =>
374 |               NewAxis (notElemAfterRemoving a.name as notElem toRemove)
375 |             ExistingAxis e prf =>
376 |               let (e' ** prfIdx) = indexAfterRemoving a.name as e toRemove
377 |                                      {neq = proofIneqIsNo contra}
378 |               in ExistingAxis e' (rewrite prfIdx in prf)
379 |     
380 | notElemExample1 : NotElem "i" ["g" ~> List, "j" ~> BinTree]
381 | notElemExample1 = %search
382 |
383 | tensorShapeTest1 : TensorShape 2
384 | tensorShapeTest1 = ["batchSize" ~> Vect 128, "seqLen" ~> List]
385 |
386 | tensorShapeTest2 : TensorShape 3
387 | tensorShapeTest2
388 |   = ["batchSize" ~> Vect 128, "seqLen" ~> List, "batchSize" ~> Vect 128]
389 |
390 | failing
391 |   tensorShapeTest3 : TensorShape 2
392 |   tensorShapeTest3 = ["batchSize" ~> Vect 128, "batchSize" ~> Vect 13]
393 |
394 | uniqueElemExample1 : UniqueElem "j" ["i" ~> List, "j" ~> BinTree, "i" ~> List]
395 | uniqueElemExample1 = %search
396 |
397 | failing
398 |   uniqueElemExampleFail : UniqueElem "x" ["i" ~> List, "j" ~> BinTree]
399 |   uniqueElemExampleFail = %search
400 |
401 |   uniqueElemExampleFail2 : UniqueElem "i" [ "i" ~> List
402 |                                           , "j" ~> BinTree
403 |                                           , "i" ~> List]
404 |   uniqueElemExampleFail2 = %search
405 |
406 | public export
407 | TensorTest1 : TensorShape 3
408 | TensorTest1 = ["batchSize" ~> Vect 128, "seqLen" ~> List, "feat" ~> Vect 64]
409 |
410 | -- Here proof search does not work (via keycommand), but `%search` does
411 | public export
412 | TensorTest2 : (i : Axis) -> ConsistentWith i [i]
413 | TensorTest2 i = %search
414 |
415 | failing
416 |   TensorElemTest2 : Elem "asdf" TensorTest1
417 |   TensorElemTest2 = %search
418 |
419 | {-
420 | public export
421 | countOccurence : AxisName -> TensorShape rank -> Nat
422 | countOccurence str [] = 0
423 | countOccurence str (a :: as) = case str == a.name of 
424 |   True => 1 + countOccurence str as
425 |   False => countOccurence str as
426 |
427 | public export
428 | countAfterRemoving : AxisName -> TensorShape rank -> Nat
429 | countAfterRemoving str [] = 0
430 | countAfterRemoving str (a :: as) with (decEq str a.name)
431 |   _ | (Yes _) = countAfterRemoving str as
432 |   _ | (No _) = 1 + countAfterRemoving str as
433 |
434 | public export
435 | countLessThanRank : (axisN : AxisName) ->
436 |   (shape : TensorShape rank) ->
437 |   LTE (countOccurence axisN shape) rank
438 | countLessThanRank axisN [] = LTEZero
439 | countLessThanRank axisN ((name ~> _) :: as) with (axisN == name)
440 |   _ | True = LTESucc (countLessThanRank axisN as)
441 |   _ | False = lteSuccRight (countLessThanRank axisN as)
442 |
443 | -- mutual
444 | --   public export
445 | --   removeAxis : (toRemove : AxisName) ->
446 | --     (shape : TensorShape rank) ->
447 | --     TensorShape (countAfterRemoving toRemove shape)
448 | --   removeAxis toRemove [] = []
449 | --   removeAxis toRemove (a :: as) with (decEq toRemove a.name)
450 | --     _ | (Yes prf) = removeAxis toRemove as
451 | --     _ | (No contra) = let cProof = consistentAfterRemoving a as toRemove
452 | --                       in a :: removeAxis toRemove as
453 | --     
454 | --   ||| Proof that, if `a` is consistent with the shape `as`, then removing any
455 | --   ||| axis from `ss` will still keep `s` consistent with the result
456 | --   public export
457 | --   consistentAfterRemoving : (a : Axis) -> (as : TensorShape rank) ->
458 | --     a `ConsistentWith` as =>
459 | --     (toRemove : AxisName) ->
460 | --     a `ConsistentWith` (removeAxis toRemove as)
461 | --   consistentAfterRemoving _ [] _ = NewAxis NotInEmpty
462 | --   consistentAfterRemoving a (ax :: as) @{(NewAxis (NotInNonEmpty neq axisConsistent))} toRemove = ?eifhh_2
463 | --   consistentAfterRemoving a (ax :: as) @{(ExistingAxis e prf)} toRemove = ?eifhh_1
464 | -}
465 |
466 |
467 |
468 |
469 | {-
470 | ||| Proof that an axis name appears in a tensor shape
471 | ||| The proof indirectly carries data of the first index of the occurence
472 | public export
473 | data InShape : AxisName -> TensorShape rank -> Type where
474 |   Here : {as : TensorShape rank} ->
475 |     (axisName ~> c) `ConsistentWith` as => -- todo add maybe inshape?
476 |     -- isno?
477 |     -- implement Elem and NotElem, and use them here?
478 |     InShape axisName ((axisName ~> c) :: as)
479 |   There : {as : TensorShape rank} -> (is : InShape axisName as) =>
480 |     a `ConsistentWith` as =>
481 |     InShape axisName (a :: as)
482 |
483 | -- ||| TODO rethink this function?
484 | -- ||| In a tensor shape removes all but the first occurence of an axis
485 | -- ||| removeDuplicates ["x" ~> 1, "y" ~> 3, "x" ~> 1] "x" = ["x" ~> 1, "y" ~> 1]
486 | -- public export
487 | -- removeDuplicates : {n, rank : Nat} -> LTE n rank =>
488 | --   (shape : TensorShape rank) ->
489 | --   (axisName : AxisName) ->
490 | --   (inShape : InShape axisName shape n) =>
491 | --   IsSucc n =>
492 | --   (m : Nat ** TensorShape m)
493 | -- removeDuplicates shape axisName {inShape} {n = 1}
494 | --   = (rank ** shape)
495 | -- removeDuplicates ((_ ~> a) :: as) axisName {inShape = Here @{is}} {n = (S (S k))}
496 | --   = removeDuplicates as axisName {inShape=is}
497 | -- removeDuplicates (s :: as) axisName {inShape = There @{is}} {n = (S (S k))}
498 | --   = let (m ** as') = removeDuplicates as axisName {inShape=is}
499 | --     in (S m ** (::) {axisConsistent=(believe_me ())} s as')
500 |
501 | -- Does tensor contraction allow duplicate axis names
502 | --   * in the input (yes, this is what Einsum also allows)
503 | --   * in the output (no, because otherwise its not clear what should happen)
504 | --     * this means that we can't write `einsum("i,i->ii")`
505 | -- 3) How does contraction work?
506 | --   3.1) Given `t : Tensor [BatchSize, BatchSize] Double`, what is `dotGeneral t`?
507 | -- 
508 | -- Need to figure out how `reduce name t` acts when:
509 | -- 1) `name="BatchSize"` and `t : Tensor [BatchSize, BatchSize] Double`
510 | --   - Should sum up the diagonal?
511 | -- 2) `name="BatchSize"` and `t : Tensor [BatchSize] Double`
512 | --   - Should sum up the vector?
513 | -- 3) `name="BatchSize"` and `t : Tensor [BatchSize, SeqLen, BatchSize] Double`
514 | --   - Should sum up the diagonal slices of SeqLen
515 | -- 
516 | -- I suppose this is about iterators
517 | -- iterating through 
518 |
519 |
520 |   -- ||| If an axis `i` can be added into a singleton list `[j]`, then
521 |   -- ||| the axis `j` can be added into a singleton list `[i]`
522 |   -- public export
523 |   -- axisConsistentSym : {i, j : Axis} ->
524 |   --   ConsistentWith i [j] -> ConsistentWith j [i]
525 |   -- axisConsistentSym (NewAxis ne) = NewAxis (notElemSym ne)
526 |   -- -- For some reason we can't pattern match on `Here`? The proof should still 
527 |   -- -- be fine... 
528 |   -- axisConsistentSym (ExistingAxis (There Here) _) impossible
529 |   -- axisConsistentSym (ExistingAxis (There (There later)) _) impossible
530 |
531 |
532 |
533 | --  public export
534 | --  data InShape : AxisName -> TensorShape rank -> (n : Nat) ->
535 | --    (ltee : LTE n rank) => Type where
536 | --    Here : {0 n, rank : Nat} -> (lte : LTE n rank) =>
537 | --      {as : TensorShape rank} -> InShape axisName as n =>
538 | --      (axisName ~> c) `ConsistentWith` as =>
539 | --      InShape {rank=S rank} axisName ((axisName ~> c) :: as) (S n) @{LTESucc lte}
540 | --    There : {0 n, rank : Nat} -> (lte : LTE n rank) =>
541 | --      {as : TensorShape rank} -> (is : InShape axisName as n) =>
542 | --      a `ConsistentWith` as =>
543 | --      InShape {rank=S rank} axisName (a :: as) n @{lteSuccRight lte}
544 | --
545 | --
546 | --  tensorShapeTest11 : TensorShape 2
547 | --  tensorShapeTest11 = ["batchSize" ~> Vect 128, "seqLen" ~> List]
548 | --
549 | --  -- ttt : (n : Nat ** InShape "batchSize" tensorShapeTest11 n @{LTEZero})
550 | --
551 | --  mutual
552 | --    public export
553 | --    removeAxis : {n, rank : Nat} -> (lte : LTE n rank) =>
554 | --      (shape : TensorShape rank) ->
555 | --      (toRemove : AxisName) ->
556 | --      (inShape : InShape {rank=rank} toRemove shape n @{lte}) =>
557 | --      TensorShape (rank `minus` n)
558 | --    removeAxis {rank = 0} shape _ = shape
559 | --    removeAxis {n=S _, rank = S _} ((toRemove ~> c) :: as) toRemove {inShape=Here @{lte} @{is}}
560 | --      = removeAxis as toRemove
561 | --    removeAxis {rank = (S rank')} (a :: as) toRemove {inShape=There @{lte} @{is}}
562 | --      = rewrite minusSuccLTE lte in
563 | --        (let consistencyProof = consistentAfterRemoving a as toRemove {is=is}
564 | --         in a :: removeAxis as toRemove)