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