0 | module Data.Tensor.Shape.Axis
  1 |
  2 | import Data.Vect.Quantifiers
  3 |
  4 | import Data.Container.Base
  5 | import Misc
  6 |
  7 | ||| The name for an axis is an arbitrary string
  8 | public export
  9 | AxisName : Type
 10 | AxisName = String
 11 |
 12 | ||| An axis is a container together with its name
 13 | public export
 14 | record Axis where
 15 |   constructor (~>)
 16 |   name : AxisName
 17 |   cont : Cont
 18 |
 19 | public export infixr 0 ~> -- Constructor for container-based axes
 20 | public export infixr 0 ~~> -- 'Constructor' for cubical axes
 21 |
 22 | public export
 23 | rename : Axis -> AxisName -> Axis
 24 | rename a str = str ~> a.cont
 25 |
 26 | ||| In some cases we TensorType might need to assign a default name to an axis,
 27 | ||| one which is internal and will not be exposed to the user.
 28 | ||| This is the default name for such cases
 29 | public export
 30 | TTInternalName : AxisName
 31 | TTInternalName = "__tensortype_tempaxis__"
 32 |
 33 | ||| Variant of container `At` for axes
 34 | public export
 35 | At : (ax : Axis) -> ax.cont.Shp -> Axis
 36 | At ax s = ax.name ~> ax.cont `At` s
 37 |
 38 | namespace Cubical
 39 |   ||| A "constructor" for cubical axes
 40 |   public export
 41 |   (~~>) : AxisName -> Nat -> Axis
 42 |   (~~>) axisName n = axisName ~> Vect n
 43 |
 44 |   ||| Follows the pattern of `IsCubical` from `Base.Object.Definition`
 45 |   public export
 46 |   data IsCubical : Axis -> Type where
 47 |     MkIsCubical : (name : AxisName) -> (n : Nat) -> IsCubical (name ~~> n)
 48 |
 49 |   ||| Evidence of axis cubicality -> evidence of underlying container cubicality
 50 |   public export
 51 |   toContCubical : {0 a : Axis} -> IsCubical a -> IsCubical a.cont
 52 |   toContCubical (MkIsCubical _ n) = Via n
 53 |
 54 |   ||| Extract the dimension from IsCubical with axis implicit
 55 |   public export
 56 |   dimHelper : {0 a : Axis} -> IsCubical a -> Nat
 57 |   dimHelper (MkIsCubical _ n) = n
 58 |
 59 |   ||| Extract the dimension from an axis which we know is cubical
 60 |   public export
 61 |   (.dim) : (0 a : Axis) -> IsCubical a => Nat
 62 |   (.dim) _ @{ic} = dimHelper ic
 63 |
 64 |   ||| Extract the dimensions of cubical axes, with shape implicit
 65 |   public export
 66 |   dimsHelper : {0 shape : Vect r Axis} ->
 67 |     All IsCubical shape -> List Nat
 68 |   dimsHelper [] = []
 69 |   dimsHelper (ic :: ns) = dimHelper ic :: dimsHelper ns
 70 |
 71 |   ||| Extract the dimensions of cubical axes, with shape explicit
 72 |   public export
 73 |   dims : (0 shape : Vect r Axis) -> All IsCubical shape => List Nat
 74 |   dims _ @{ac} = dimsHelper ac
 75 |
 76 |   ||| Product of all the dimensions of a cubical tensors, i.e. its size
 77 |   public export
 78 |   size : (0 shape : Vect r Axis) -> (ac : All IsCubical shape) => Nat
 79 |   size shape = prod (dims shape)
 80 |
 81 | namespace Naperian
 82 |   ||| Follows the pattern of `IsNaperian` from `Base.Object.Definition`
 83 |   public export
 84 |   data IsNaperian : Axis -> Type where
 85 |     MkIsNaperian : (name : AxisName) -> (pos : Type) ->
 86 |       IsNaperian (name ~> Nap pos)
 87 |
 88 |   ||| Evidence of axis being Naperian -> evidence of container being Naperian
 89 |   %hint
 90 |   public export
 91 |   toContNaperian : {0 a : Axis} -> IsNaperian a -> IsNaperian a.cont
 92 |   toContNaperian (MkIsNaperian _ pos) = Via pos
 93 |   
 94 |   ||| Extract the position type from IsNaperian with axis implicit
 95 |   public export
 96 |   LogHelper : {0 a : Axis} -> IsNaperian a => Type
 97 |   LogHelper @{MkIsNaperian _ pos} = pos
 98 |
 99 |   ||| Extract the position type from an axis which we know is Naperian
100 |   public export
101 |   Log : (0 a : Axis) -> IsNaperian a => Type
102 |   Log a @{inn} = LogHelper @{inn}
103 |
104 | namespace IsConcrete
105 |   public export
106 |   data IsConcrete : Axis -> Type where
107 |     MkIsConcrete : (name : AxisName) -> IsConcrete c -> IsConcrete (name ~> c)
108 |
109 |   public export
110 |   toContConcrete : {0 a : Axis} -> IsConcrete a -> IsConcrete a.cont
111 |   toContConcrete (MkIsConcrete _ ic) = ic