0 | {--
  1 | Copyright (C) 2022  Joel Berkeley
  2 |
  3 | This program is free software: you can redistribute it and/or modify
  4 | it under the terms of the GNU Affero General Public License as published
  5 | by the Free Software Foundation, either version 3 of the License, or
  6 | (at your option) any later version.
  7 |
  8 | This program is distributed in the hope that it will be useful,
  9 | but WITHOUT ANY WARRANTY; without even the implied warranty of
 10 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 11 | GNU Affero General Public License for more details.
 12 |
 13 | You should have received a copy of the GNU Affero General Public License
 14 | along with this program.  If not, see <https://www.gnu.org/licenses/>.
 15 | --}
 16 | ||| Defines `Literal`, a single value or array of values with a specified shape.
 17 | ||| `Literal` is similar to `Tensor`, but differs in a number of important ways:
 18 | |||
 19 | ||| * `Literal` offers a convenient syntax for constructing `Literal`s with boolean and numeric
 20 | |||   contents. For example, `True`, `1` and `[1, 2, 3]` are all valid `Literal`s. This makes it
 21 | |||   useful for constructing `Tensor`s.
 22 | ||| * Operations on `Literal` are *not* accelerated by a graph compiler, so operations on large
 23 | |||   `Literal`s, and large sequences of operations on any `Literal`, can be expected to be slower
 24 | |||   than they would on an equivalent `Tensor`.
 25 | ||| * `Literal` is implemented in pure Idris. As such, values can contain elements of any type, and
 26 | |||   implements a number of standard Idris interfaces. This, along with its convenient syntax,
 27 | |||   makes it particularly useful for testing operations on `Tensor`s.
 28 | module Literal
 29 |
 30 | import Data.List.Elem
 31 | import public Types
 32 |
 33 | ||| A scalar or array of values.
 34 | public export
 35 | data Literal : Shape -> Type -> Type where
 36 |   Scalar : a -> Literal [] a
 37 |   Nil : Literal (0 :: ds) a
 38 |   (::) : Literal ds a -> Literal (d :: ds) a -> Literal (S d :: ds) a
 39 |
 40 | export
 41 | fromInteger : Num a => Integer -> Literal [] a
 42 | fromInteger = Scalar . fromInteger
 43 |
 44 | export
 45 | fromDouble : Double -> Literal [] Double
 46 | fromDouble = Scalar
 47 |
 48 | ||| Convenience aliases for scalar boolean literals.
 49 | export
 50 | True, False : Literal [] Bool
 51 | True = Scalar True
 52 | False = Scalar False
 53 |
 54 | export
 55 | Functor (Literal shape) where
 56 |   map f (Scalar x) = Scalar (f x)
 57 |   map _ [] = []
 58 |   map f (x :: xs) = map f x :: map f xs
 59 |
 60 | functorIdentity : (xs : Literal shape a) -> map Prelude.id xs = xs
 61 | functorIdentity (Scalar _) = Refl
 62 | functorIdentity [] = Refl
 63 | functorIdentity (x :: xs) = cong2 (::) (functorIdentity x) (functorIdentity xs)
 64 |
 65 | functorComposition :
 66 |   (xs : Literal shape a) -> (f : a -> b) -> (g : b -> c) -> map (g . f) xs = map g (map f xs)
 67 | functorComposition (Scalar _) _ _ = Refl
 68 | functorComposition [] _ _ = Refl
 69 | functorComposition (x :: xs) f g =
 70 |   cong2 (::) (functorComposition x f g) (functorComposition xs f g)
 71 |
 72 | export
 73 | {shape : _} -> Applicative (Literal shape) where
 74 |   pure x = case shape of
 75 |     [] => Scalar x
 76 |     (0 :: _) => []
 77 |     (S d :: ds) => pure x :: assert_total (pure x)
 78 |
 79 |   (Scalar f) <*> (Scalar x) = Scalar (f x)
 80 |   [] <*> [] = []
 81 |   (f :: fs) <*> (x :: xs) = (f <*> x) :: (fs <*> xs)
 82 |
 83 | applicativeIdentity : (xs : Literal shape a) -> pure Prelude.id <*> xs = xs
 84 | applicativeIdentity (Scalar _) = Refl
 85 | applicativeIdentity [] = Refl
 86 | applicativeIdentity (x :: xs) = cong2 (::) (applicativeIdentity x) (applicativeIdentity xs)
 87 |
 88 | applicativeHomomorphism :
 89 |   (shape : Shape) -> (f : a -> b) -> (x : a) -> pure f <*> pure x = pure {f = Literal shape} (f x)
 90 | applicativeHomomorphism [] _ _ = Refl
 91 | applicativeHomomorphism (d :: ds) f x = forCons d ds f x
 92 |   where
 93 |   forCons :
 94 |     (d : Nat) ->
 95 |     (ds : Shape) ->
 96 |     (f : a -> b) ->
 97 |     (x : a) ->
 98 |     pure f <*> pure x = pure {f = Literal (d :: ds)} (f x)
 99 |   forCons 0 _ _ _ = Refl
100 |   forCons (S d) ds f x = cong2 (::) (applicativeHomomorphism ds f x) (forCons d ds f x)
101 |
102 | applicativeInterchange :
103 |   (fs : Literal shape (a -> b)) -> (x : a) -> fs <*> pure x = pure (x) <*> fs
104 | applicativeInterchange (Scalar _) _ = Refl
105 | applicativeInterchange [] _ = Refl
106 | applicativeInterchange (f :: fs) x =
107 |   cong2 (::) (applicativeInterchange f x) (applicativeInterchange fs x)
108 |
109 | applicativeComposition :
110 |   (xs : Literal shape a) ->
111 |   (fs : Literal shape (a -> b)) ->
112 |   (gs : Literal shape (b -> c)) ->
113 |   pure (.) <*> gs <*> fs <*> xs = gs <*> (fs <*> xs)
114 | applicativeComposition (Scalar _) (Scalar _) (Scalar _) = Refl
115 | applicativeComposition [] [] [] = Refl
116 | applicativeComposition (x :: xs) (f :: fs) (g :: gs) =
117 |   cong2 (::) (applicativeComposition x f g) (applicativeComposition xs fs gs)
118 |
119 | export
120 | Foldable (Literal shape) where
121 |   foldr f acc (Scalar x) = f x acc
122 |   foldr _ acc [] = acc
123 |   foldr f acc (x :: xs) = foldr f (foldr f acc xs) x
124 |
125 | export
126 | Traversable (Literal shape) where
127 |   traverse f (Scalar x) = [| Scalar (f x) |]
128 |   traverse f [] = pure []
129 |   traverse f (x :: xs) = [| traverse f x :: traverse f xs |]
130 |
131 | export
132 | Zippable (Literal shape) where
133 |   zipWith f (Scalar x) (Scalar y) = Scalar (f x y)
134 |   zipWith _ [] [] = []
135 |   zipWith f (x :: xs) (y :: ys) = zipWith f x y :: zipWith f xs ys
136 |
137 |   zipWith3 f (Scalar x) (Scalar y) (Scalar z) = Scalar (f x y z)
138 |   zipWith3 _ [] [] [] = []
139 |   zipWith3 f (x :: xs) (y :: ys) (z :: zs) = zipWith3 f x y z :: zipWith3 f xs ys zs
140 |
141 |   unzipWith f (Scalar x) = let (x, y) = f x in (Scalar x, Scalar y)
142 |   unzipWith _ [] = ([], [])
143 |   unzipWith f (x :: xs) =
144 |     let (x, y) = unzipWith f x
145 |         (xs, ys) = unzipWith f xs
146 |      in (x :: xs, y :: ys)
147 |
148 |   unzipWith3 f (Scalar x) = let (x, y, z) = f x in (Scalar x, Scalar y, Scalar z)
149 |   unzipWith3 _ [] = ([], [], [])
150 |   unzipWith3 f (x :: xs) =
151 |     let (x, y, z) = unzipWith3 f x
152 |         (xs, ys, zs) = unzipWith3 f xs
153 |      in (x :: xs, y :: ys, z :: zs)
154 |
155 | ||| `True` if no elements are `False`. `all []` is `True`.
156 | export
157 | all : Literal shape Bool -> Bool
158 | all xs = foldr (\x, y => x && y) True xs
159 |
160 | export
161 | Num a => Num (Literal [] a) where
162 |   x + y = [| x + y |]
163 |   x * y = [| x * y |]
164 |   fromInteger = Scalar . fromInteger
165 |
166 | export
167 | negate : Neg a => Literal shape a -> Literal shape a
168 | negate = map negate
169 |
170 | export
171 | Eq a => Eq (Literal shape a) where
172 |   x == y = all (zipWith (==) x y)
173 |
174 | toVect : Literal (d :: ds) a -> Vect d (Literal ds a)
175 | toVect [] = []
176 | toVect (x :: y) = x :: toVect y
177 |
178 | ||| Show the `Literal`. The `Scalar` constructor is omitted for brevity.
179 | export
180 | {shape : _} -> Show a => Show (Literal shape a) where
181 |   show = showWithIndent "" where
182 |     showWithIndent : {shape : _} -> String -> Literal shape a -> String
183 |     showWithIndent _ (Scalar x) = show x
184 |     showWithIndent _ [] = "[]"
185 |     showWithIndent {shape = [S _]} _ x = show (toList x)
186 |     showWithIndent {shape = (S d :: dd :: ddd)} indent (x :: xs) =
187 |       let indent = " " ++ indent
188 |           first = showWithIndent indent x
189 |           rest = foldMap (\e => ",\n" ++ indent ++ showWithIndent indent e) (toVect xs)
190 |        in "[" ++ first ++ rest ++ "]"
191 |
192 | export
193 | {shape : _} -> Cast (Array shape a) (Literal shape a) where
194 |   cast x with (shape)
195 |     cast x | [] = Scalar x
196 |     cast _ | (0 :: _) = []
197 |     cast (x :: xs) | (S d :: ds) = cast x :: cast xs
198 |
199 | export
200 | [toArray] Cast (Literal shape a) (Array shape a) where
201 |   cast (Scalar x) = x
202 |   cast [] = []
203 |   cast (x :: y) = cast @{toArray} x :: cast @{toArray} y
204 |
205 | namespace All
206 |   ||| An `All p xs` is an array (or scalar) of proofs about each element in `xs`.
207 |   |||
208 |   ||| For example, an `All IsSucc xs` proves that every element in `xs` is non-zero.
209 |   public export
210 |   data All : (0 p : a -> Type) -> Literal shape a -> Type where
211 |     Scalar : forall x . p x -> All p (Scalar x)
212 |     Nil  : All p []
213 |     (::) : All p x -> All p xs -> All p (x :: xs)
214 |
215 | namespace All2
216 |   ||| An `All2 p xs ys` is an array (or scalar) of pairwise proofs about elements in `xs` and `ys`.
217 |   |||
218 |   ||| For example, an `All2 LT xs ys` proves that each number in `xs` is less than the number in
219 |   ||| `ys` at the same position.
220 |   public export
221 |   data All2 : (p : a -> b -> Type) -> Literal shape a -> Literal shape b -> Type where
222 |     Scalar : forall a, b . p a b -> All2 p (Scalar a) (Scalar b)
223 |     Nil : All2 p [] []
224 |     (::) : All2 p a b -> All2 p as bs -> All2 p (a :: as) (b :: bs)
225 |