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