
   0 | module Data.CT.DependentPara.Instances
   1 | 
   2 | import Data.DPair
   3 | import Data.CT.Category.Definition
   4 | import Data.CT.Functor.Definition
   5 | import Data.CT.DependentAction.Definition
   6 | import Data.CT.DependentPara.Definition
   7 | import Data.CT.Category.Instances
   8 | import Data.CT.Functor.Instances
   9 | import Data.CT.DependentAction.Instances
  10 | 
  11 | import Data.Container.Base
  12 | import Data.Container.Additive
  13 | 
  14 | {-------------------------------------------------------------------------------
  15 | {-------------------------------------------------------------------------------
  16 | Default Para is dependent Para, and default lenses are dependent additive lenses
  17 | 
  18 | Ideally, this code would uphold the notation principle that
  19 | `-\->` denotes dependent parametric functions and
  20 | `=\\=>` denotes dependent parametric lenses
  21 | 
  22 | But since learning where the parameter depends on the input hasn't been
  23 | implemented yet, the code will write `=\\=>` for non-dependent additive lenses, and simply not export an infix notation for dependent additive lenses
  24 | 
  25 | Likewise, instead going in and defining full-blown definitions of dependent 
  26 | actegories with units and coherences we instead leverage the main definition in 
  27 | the background and only instantiate cases, manually:
  28 | one for parametric functions and one for parametric additive dependent lenses.
  29 | We instantiate them using same names in different namespaces, and leverage Idris' name resolution mechanisms to allow the user to use the same name and
  30 | reduce cognitive load.
  31 | However, to reduce typechecking time, there are also now some concrete records.
  32 | -------------------------------------------------------------------------------}
  33 | -------------------------------------------------------------------------------}
  34 | 
  35 | -- Para(Set)
  36 | public export infixr 1 -\-> -- dependent parametric functions
  37 | public export infixr 1 -\--> -- non-dependent parametric functions
  38 | 
  39 | -- Para(AddCont)
  40 | public export infixr 1 =\\=> -- non-dependent parametric lenses
  41 | -- public export infixr 1 =\\==> -- not exported, see comment at top of file
  42 | 
  43 |   
  44 | -- Dependent parametric function composition
  45 | public export infixr 10 \>>
  46 | -- Dependent parametric lens composition
  47 | public export infixr 10 &>>
  48 | 
  49 | ||| The 2-category `Para(Set)`
  50 | namespace ParametricFunctions
  51 |   ||| Non-dependent parametric functions
  52 |   public export
  53 |   Para : (a, b : Type) -> Type
  54 |   Para = DepParaMor PairType
  55 | 
  56 |   ||| Infix notation for non-dependent parametric functions
  57 |   ||| We interpret the extra "-" as a mental symbol for "flat",
  58 |   ||| i.e. "non-dependent"
  59 |   public export
  60 |   (-\-->) : (a, b : Type) -> Type
  61 |   a -\--> b = Para a b
  62 | 
  63 |   ||| Dependent parametric functions, i.e. where the parameter type varies with
  64 |   ||| values of the input `a`
  65 |   public export
  66 |   DPara : (a, b : Type) -> Type 
  67 |   DPara = DepParaMor DPairType
  68 |   
  69 |   ||| Infix notation for dependent parametric functions
  70 |   ||| We interpret the crossed line as a parameter coming in from the top
  71 |   public export
  72 |   (-\->) : (a, b : Type) -> Type
  73 |   a -\-> b = DPara a b
  74 | 
  75 |   public export
  76 |   trivialParam : (a -> b) -> a -\-> b
  77 |   trivialParam f = MkPara 
  78 |     (\_ => Unit)
  79 |     (\(a ** ()) => f a)
  80 | 
  81 |   public export
  82 |   id : a -\-> a
  83 |   id = trivialParam id
  84 | 
  85 |   public export
  86 |   composePara : a -\-> b -> b -\-> c -> a -\-> c
  87 |   composePara (MkPara p f) (MkPara q g) = MkPara
  88 |     (\x => DPair (p x) (\p' => q (f (x ** p'))) )
  89 |     (\(x ** (p' ** q')) => g (f (x ** p') ** q'))
  90 | 
  91 |   public export
  92 |   composeParallel : a -\-> b -> c -\-> d -> (a, c) -\-> (b, d)
  93 |   composeParallel (MkPara p f) (MkPara q g) = MkPara
  94 |     (\(x, y) => (p x, q y))
  95 |     (\((x, y) ** (px, qy)) => (f (x ** px), g (y ** qy)))
  96 |   
  97 |   public export
  98 |   (\>>) : a -\-> b -> b -\-> c -> a -\-> c
  99 |   (\>>) = composePara
 100 | 
 101 |   public export
 102 |   reparam : (pf : a -\-> b) ->
 103 |     {q : a -> Type} ->
 104 |     (r : (x : a) -> q x -> pf.Param x) ->
 105 |     a -\-> b
 106 |   reparam (MkPara p f) r = MkPara q (\(x ** qq) => f (x ** (r x qq)))
 107 | 
 108 |   public export
 109 |   Param : DPara a b -> a -> Type
 110 |   Param = DepParaMor.Param
 111 |   
 112 |   public export
 113 |   Run : (pf : DPara a b) -> (x : a) -> Param pf x -> b
 114 |   Run pf = DPair.curry (DepParaMor.Run pf)
 115 | 
 116 |   public export
 117 |   data IsNotDependent : DPara a b -> Type where
 118 |     MkNonDep : (p : Type) -> (f : DPair a (const p) -> b) ->
 119 |       IsNotDependent (MkPara (\_ => p) f)
 120 |   
 121 |   public export
 122 |   GetNonDep : (pf : DPara a b) ->
 123 |     IsNotDependent pf => (p : Type ** DPair a (const p) -> b)
 124 |   GetNonDep _ @{MkNonDep p f} = (p ** f)
 125 | 
 126 |   ||| Get the parameter of a non-dependent parametric function
 127 |   public export
 128 |   GetParam : (pf : DPara a b) ->
 129 |     IsNotDependent pf => Type
 130 |   GetParam _ @{MkNonDep p f} = p
 131 | 
 132 |   public export
 133 |   composeNTimes : Nat -> a -\-> a -> a -\-> a
 134 |   composeNTimes 0 f = id
 135 |   composeNTimes 1 f = f -- to get rid of the annoying Unit parameter
 136 |   composeNTimes (S k) f = composePara f (composeNTimes k f)
 137 | 
 138 |   public export
 139 |   binaryOpToPara : {p : Type} -> (f : (a, p) -> b) -> a -\-> b
 140 |   binaryOpToPara f = MkPara
 141 |     (\_ => p)
 142 |     (\(x ** p') => f (x, p'))
 143 | 
 144 | ||| The 2-category Para(AddCont)
 145 | ||| Para(Cont) is not used in TensorType
 146 | namespace ParametricLenses
 147 |   ||| Non-dependent parametric lenses.
 148 |   ||| As mentioned on top, all of these lenses are additive, and dependent
 149 |   |||
 150 |   ||| As a record, because otherwise the implicit argument carrying slows down
 151 |   ||| typechecking practical neural network architectures. That is, every
 152 |   ||| `.Param` and `.Run` carry `AddDLens`, `Const` and `PairAddCont` as 
 153 |   ||| implcit arguments, unfolded into the full `MkCat` and `MkFunctor` 
 154 |   ||| structure. This happens in bodies of, say `>*<`, at *every occurence*
 155 |   public export
 156 |   record ParaAddLens (a, b : AddCont) where
 157 |     constructor MkPara
 158 |     Param : AddCont
 159 |     Run : (a >*< Param) =%+> b
 160 | 
 161 |   ||| Infix notation for non-dependent parametric additive lenses
 162 |   ||| Compared to `-\-->`, every line is doubled, meant to be interpreted as 
 163 |   ||| information flowing bidirectionally. 
 164 |   ||| See comment for top of file for further explanation
 165 |   public export
 166 |   (=\\=>) : (a, b : AddCont) -> Type
 167 |   a =\\=> b = ParaAddLens a b
 168 | 
 169 |   namespace Wrapping
 170 |     ||| Simple wrapping and unwrapping because this is a record now
 171 |     public export
 172 |     toDepPara : ParaAddLens a b -> DepParaMor PairAddCont a b
 173 |     toDepPara (MkPara p f) = MkPara p f
 174 | 
 175 |     public export
 176 |     fromDepPara : DepParaMor PairAddCont a b -> ParaAddLens a b
 177 |     fromDepPara (MkPara p f) = MkPara p f
 178 | 
 179 | 
 180 |   public export
 181 |   trivialParam : a =%+> b -> a =\\=> b
 182 |   trivialParam f = MkPara
 183 |     UnitCont
 184 |     (!%+ \(x, ()) =>
 185 |       let (y ** ky) = (%!+) f x
 186 |       in (y ** \y' => (ky y', ())))
 187 | 
 188 |   public export
 189 |   binaryOpToPara : {p : AddCont} ->
 190 |     (a >*< p) =%+> b -> a =\\=> b
 191 |   binaryOpToPara f = MkPara p f
 192 | 
 193 |   public export
 194 |   id : a =\\=> a
 195 |   id = trivialParam id
 196 | 
 197 |   public export
 198 |   toHomRepresentation : (f : ParaAddLens a b) ->
 199 |     (Param f) =%+> InternalLensAdditive a b
 200 |   toHomRepresentation (MkPara pType f) = !%+ \p =>
 201 |     (!%+ \a => (f.fwd (a, p) ** \b' => fst (f.bwd (a, p) b')) **
 202 |       \l => foldr (\(a ** b') => pType.Plus p (snd (f.bwd (a, p) b'))) (pType.Zero p) l)
 203 | 
 204 |   public export
 205 |   composePara : a =\\=> b -> b =\\=> c -> a =\\=> c
 206 |   composePara f g = MkPara
 207 |     (Param f >*< Param g)
 208 |     (assocR %+>> (Run f >*< id) %+>> Run g)
 209 | 
 210 |   public export
 211 |   composeParallel : a =\\=> b -> c =\\=> d -> (a >*< c) =\\=> (b >*< d)
 212 |   composeParallel f g = MkPara
 213 |     (Param f >*< Param g)
 214 |     (swapMiddle %+>> (Run f >*< Run g))
 215 | 
 216 |   ||| Postcompose with a lens
 217 |   public export
 218 |   postcomposeLens : a =\\=> b -> b =%+> c -> a =\\=> c
 219 |   postcomposeLens f g = MkPara (Param f) (Run f %+>> g)
 220 | 
 221 | 
 222 | namespace DependentParametricLenses
 223 |   ||| Dependent parametric lenses, i.e. where the parameter container can vary
 224 |   ||| with the shape of the input container
 225 |   ||| Defined as its own record for the same reason as `ParaAddLens`
 226 |   public export
 227 |   record DParaAddLens (a, b : AddCont) where
 228 |     constructor MkPara
 229 |     Param : a.Shp -> AddCont
 230 |     Run : DPair a Param =%+> b
 231 | 
 232 |   namespace Wrap
 233 |     public export
 234 |     toDepPara : DParaAddLens a b -> DepParaMor DPairAddCont a b
 235 |     toDepPara (MkPara p f) = MkPara p f
 236 | 
 237 |     public export
 238 |     fromDepPara : DepParaMor DPairAddCont a b -> DParaAddLens a b
 239 |     fromDepPara (MkPara p f) = MkPara p f
 240 | 
 241 |   {- commented out for now, since its not used
 242 |   ||| Infix notation for additive parametric dependent lenses
 243 |   public export
 244 |   (=\\=>) : (a, b : AddCont) -> Type
 245 |   a =\\=> b = DParaAddLens a b
 246 |   
 247 |   public export
 248 |   trivialParam : a =%+> b -> a =\\=> b
 249 |   trivialParam f = MkPara
 250 |     (\_ => UnitCont)
 251 |     (!% !% \(x ** ()) => let (y ** ky) = (%!+) f x
 252 |                          in (y ** \y' => (ky y', ())))
 253 | 
 254 |   public export
 255 |   id : a =\\=> a
 256 |   id = trivialParam id
 257 |   
 258 |   public export
 259 |   composePara : a =\\=> b -> b =\\=> c -> a =\\=> c
 260 |   composePara (MkPara p f) (MkPara q g) = MkPara
 261 |     (\x => DPair (p x) (\ps => q (f.fwd (x ** ps))))
 262 |     (!%+ \(x ** (ps ** qs)) =>
 263 |       (g.fwd (f.fwd (x ** ps) ** qs) ** \cPos =>
 264 |         let (bPos, qPos) = g.bwd (f.fwd (x ** ps) ** qs) cPos
 265 |             (aPos, pPos) = f.bwd (x ** ps) bPos
 266 |         in (aPos, (pPos, qPos))))
 267 | 
 268 | 
 269 |   public export
 270 |   (&>>) : a =\\=> b -> b =\\=> c -> a =\\=> c
 271 |   (&>>) = composePara
 272 | 
 273 |   ||| A predicate witnessing that a parametric additive dependent lens has
 274 |   ||| a non-dependent (constant) parameter.
 275 |   public export
 276 |   data IsNotDependent : DParaAddLens a b -> Type where
 277 |     MkNonDep : (p : AddCont) -> (f : DPair a (const p) =%+> b) ->
 278 |       IsNotDependent {a=a} {b=b} (MkPara (\_ => p) f)
 279 |   
 280 |   public export
 281 |   GetNonDep : (pf : DParaAddLens a b) ->
 282 |     IsNotDependent pf => (pc : AddCont ** DPair a (const pc) =%+> b)
 283 |   GetNonDep _ @{MkNonDep pc f} = (pc ** f)
 284 | 
 285 |   public export
 286 |   GetParam : (pf : DParaAddLens a b) ->
 287 |     IsNotDependent pf => AddCont
 288 |   GetParam (MkPara (const p) f) @{MkNonDep p f} = p
 289 | 
 290 |   public export
 291 |   toHomRepresentation : (pf : DParaAddLens a b) ->
 292 |     IsNotDependent pf =>
 293 |     GetParam pf =%+> (InternalLensAdditive a b)
 294 |   toHomRepresentation (MkPara (const pc) f) @{MkNonDep pc f}
 295 |     = !%+ \p => (!%+ \x => (f.fwd (x ** p) ** \b' => fst (f.bwd (x ** p) b')) ** \l => foldr (\(x ** b') => pc.Plus p (snd (f.bwd (x ** p) b'))) (pc.Zero p) l)
 296 | 
 297 |   public export
 298 |   composeNTimes : Nat -> a =\\=> a -> a =\\=> a
 299 |   composeNTimes 0 f = id
 300 |   composeNTimes 1 f = f -- to get rid of the annoying Unit parameter
 301 |   composeNTimes (S k) f = composePara f (composeNTimes k f)
 302 | 
 303 |   ||| Convert a morphism from product container to one from DPair
 304 |   ||| This witnesses the isomorphism (a >< p) ≅ DPair a (const p)
 305 |   public export
 306 |   fromNonDepProduct : (a >*< p) =%+> b -> DPair a (const p) =%+> b
 307 |   fromNonDepProduct f = !%+ \(x ** p') => (%!+) f (x, p')
 308 | 
 309 | 
 310 |   %hide Data.Container.Base.Morphism.Definition.DependentLenses.(=%>)
 311 |   -}
 312 | 
 313 | -- public export
 314 | -- dependentMap : {t : a -> Type} -> (f : (x : a) -> t x) ->
 315 | --   Vect n a -> Vect n (x : a ** t x)
 316 | -- dependentMap f [] = []
 317 | -- dependentMap f (x :: xs) = (x ** f x) :: dependentMap f xs
 318 | -- 
 319 | -- public export infixr 10 <$^>
 320 | -- public export
 321 | -- (<$^>) : {t : a -> Type} -> (f : (x : a) -> t x) ->
 322 | --   Vect n a -> Vect n (x : a ** t x)
 323 | -- (<$^>) f xs = dependentMap f xs
 324 | 
 325 | 
 326 | -- composePara_rhs_1 : (p : Vect n Type) -> (q : Vect m Type)
 327 | --   -> (a -> All Prelude.id p -> b)
 328 | --   -> (b -> All Prelude.id q -> c)
 329 | --   -> (a -> All Prelude.id (p ++ q) -> c)
 330 | -- composePara_rhs_1 [] [] f g a [] = ?composePara_rhs_1_rhs_2
 331 | -- composePara_rhs_1 [] (q :: ws) f g a (pq :: pqs) = ?composePara_rhs_1_rhs_3
 332 | -- composePara_rhs_1 (p :: ps) q f g a pq = ?composePara_rhs_1_rhs_1
 333 | -- 
 334 | -- composePara : Para a n b -> Para b m c -> Para a (n + m) c
 335 | -- composePara (MkPara p f) (MkPara q g) = MkPara (p ++ q) (composePara_rhs_1 p q f g)