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 Spidr.Data.Literal
 29 |
 30 | import Data.List.Elem
 31 |
 32 | import public Spidr.Shape
 33 |
 34 | ||| A scalar or array of values.
 35 | public export
 36 | data Literal : Shape -> Type -> Type where
 37 |   Scalar : a -> Literal [] a
 38 |   Nil : Literal (0 :: ds) a
 39 |   (::) : Literal ds a -> Literal (d :: ds) a -> Literal (S d :: ds) a
 40 |
 41 | export
 42 | fromInteger : Num a => Integer -> Literal [] a
 43 | fromInteger = Scalar . fromInteger
 44 |
 45 | export
 46 | fromDouble : Double -> Literal [] Double
 47 | fromDouble = Scalar
 48 |
 49 | ||| Convenience aliases for scalar boolean literals.
 50 | export
 51 | True, False : Literal [] Bool
 52 | True = Scalar True
 53 | False = Scalar False
 54 |
 55 | export
 56 | Functor (Literal shape) where
 57 |   map f (Scalar x) = Scalar (f x)
 58 |   map _ [] = []
 59 |   map f (x :: xs) = map f x :: map f xs
 60 |
 61 | functorIdentity : (xs : Literal shape a) -> map Prelude.id xs = xs
 62 | functorIdentity (Scalar _) = Refl
 63 | functorIdentity [] = Refl
 64 | functorIdentity (x :: xs) = cong2 (::) (functorIdentity x) (functorIdentity xs)
 65 |
 66 | functorComposition :
 67 |   (xs : Literal shape a) -> (f : a -> b) -> (g : b -> c) -> map (g . f) xs = map g (map f xs)
 68 | functorComposition (Scalar _) _ _ = Refl
 69 | functorComposition [] _ _ = Refl
 70 | functorComposition (x :: xs) f g =
 71 |   cong2 (::) (functorComposition x f g) (functorComposition xs f g)
 72 |
 73 | export
 74 | {shape : _} -> Applicative (Literal shape) where
 75 |   pure x = case shape of
 76 |     [] => Scalar x
 77 |     (0 :: _) => []
 78 |     (S d :: ds) => pure x :: assert_total (pure x)
 79 |
 80 |   (Scalar f) <*> (Scalar x) = Scalar (f x)
 81 |   [] <*> [] = []
 82 |   (f :: fs) <*> (x :: xs) = (f <*> x) :: (fs <*> xs)
 83 |
 84 | applicativeIdentity : (xs : Literal shape a) -> pure Prelude.id <*> xs = xs
 85 | applicativeIdentity (Scalar _) = Refl
 86 | applicativeIdentity [] = Refl
 87 | applicativeIdentity (x :: xs) = cong2 (::) (applicativeIdentity x) (applicativeIdentity xs)
 88 |
 89 | applicativeHomomorphism :
 90 |   (shape : Shape) -> (f : a -> b) -> (x : a) -> pure f <*> pure x = pure {f = Literal shape} (f x)
 91 | applicativeHomomorphism [] _ _ = Refl
 92 | applicativeHomomorphism (d :: ds) f x = forCons d ds f x
 93 |   where
 94 |   forCons :
 95 |     (d : Nat) ->
 96 |     (ds : Shape) ->
 97 |     (f : a -> b) ->
 98 |     (x : a) ->
 99 |     pure f <*> pure x = pure {f = Literal (d :: ds)} (f x)
100 |   forCons 0 _ _ _ = Refl
101 |   forCons (S d) ds f x = cong2 (::) (applicativeHomomorphism ds f x) (forCons d ds f x)
102 |
103 | applicativeInterchange :
104 |   (fs : Literal shape (a -> b)) -> (x : a) -> fs <*> pure x = pure (x) <*> fs
105 | applicativeInterchange (Scalar _) _ = Refl
106 | applicativeInterchange [] _ = Refl
107 | applicativeInterchange (f :: fs) x =
108 |   cong2 (::) (applicativeInterchange f x) (applicativeInterchange fs x)
109 |
110 | applicativeComposition :
111 |   (xs : Literal shape a) ->
112 |   (fs : Literal shape (a -> b)) ->
113 |   (gs : Literal shape (b -> c)) ->
114 |   pure (.) <*> gs <*> fs <*> xs = gs <*> (fs <*> xs)
115 | applicativeComposition (Scalar _) (Scalar _) (Scalar _) = Refl
116 | applicativeComposition [] [] [] = Refl
117 | applicativeComposition (x :: xs) (f :: fs) (g :: gs) =
118 |   cong2 (::) (applicativeComposition x f g) (applicativeComposition xs fs gs)
119 |
120 | export
121 | Foldable (Literal shape) where
122 |   foldr f acc (Scalar x) = f x acc
123 |   foldr _ acc [] = acc
124 |   foldr f acc (x :: xs) = foldr f (foldr f acc xs) x
125 |
126 | export
127 | Traversable (Literal shape) where
128 |   traverse f (Scalar x) = [| Scalar (f x) |]
129 |   traverse f [] = pure []
130 |   traverse f (x :: xs) = [| traverse f x :: traverse f xs |]
131 |
132 | export
133 | Zippable (Literal shape) where
134 |   zipWith f (Scalar x) (Scalar y) = Scalar (f x y)
135 |   zipWith _ [] [] = []
136 |   zipWith f (x :: xs) (y :: ys) = zipWith f x y :: zipWith f xs ys
137 |
138 |   zipWith3 f (Scalar x) (Scalar y) (Scalar z) = Scalar (f x y z)
139 |   zipWith3 _ [] [] [] = []
140 |   zipWith3 f (x :: xs) (y :: ys) (z :: zs) = zipWith3 f x y z :: zipWith3 f xs ys zs
141 |
142 |   unzipWith f (Scalar x) = let (x, y) = f x in (Scalar x, Scalar y)
143 |   unzipWith _ [] = ([], [])
144 |   unzipWith f (x :: xs) =
145 |     let (x, y) = unzipWith f x
146 |         (xs, ys) = unzipWith f xs
147 |      in (x :: xs, y :: ys)
148 |
149 |   unzipWith3 f (Scalar x) = let (x, y, z) = f x in (Scalar x, Scalar y, Scalar z)
150 |   unzipWith3 _ [] = ([], [], [])
151 |   unzipWith3 f (x :: xs) =
152 |     let (x, y, z) = unzipWith3 f x
153 |         (xs, ys, zs) = unzipWith3 f xs
154 |      in (x :: xs, y :: ys, z :: zs)
155 |
156 | ||| `True` if no elements are `False`. `all []` is `True`.
157 | export
158 | all : Literal shape Bool -> Bool
159 | all xs = foldr (\x, y => x && y) True xs
160 |
161 | export
162 | Num a => Num (Literal [] a) where
163 |   x + y = [| x + y |]
164 |   x * y = [| x * y |]
165 |   fromInteger = Scalar . fromInteger
166 |
167 | export
168 | negate : Neg a => Literal shape a -> Literal shape a
169 | negate = map negate
170 |
171 | export
172 | Eq a => Eq (Literal shape a) where
173 |   x == y = all (zipWith (==) x y)
174 |
175 | toVect : Literal (d :: ds) a -> Vect d (Literal ds a)
176 | toVect [] = []
177 | toVect (x :: y) = x :: toVect y
178 |
179 | ||| Show the `Literal`. The `Scalar` constructor is omitted for brevity.
180 | export
181 | {shape : _} -> Show a => Show (Literal shape a) where
182 |   show = showWithIndent "" where
183 |     showWithIndent : {shape : _} -> String -> Literal shape a -> String
184 |     showWithIndent _ (Scalar x) = show x
185 |     showWithIndent _ [] = "[]"
186 |     showWithIndent {shape = [S _]} _ x = show (toList x)
187 |     showWithIndent {shape = (S d :: dd :: ddd)} indent (x :: xs) =
188 |       let indent = " " ++ indent
189 |           first = showWithIndent indent x
190 |           rest = foldMap (\e => ",\n" ++ indent ++ showWithIndent indent e) (toVect xs)
191 |        in "[" ++ first ++ rest ++ "]"
192 |
193 | export
194 | {shape : _} -> Cast (Array shape a) (Literal shape a) where
195 |   cast x with (shape)
196 |     cast x | [] = Scalar x
197 |     cast _ | (0 :: _) = []
198 |     cast (x :: xs) | (S d :: ds) = cast x :: cast xs
199 |
200 | export
201 | [toArray] Cast (Literal shape a) (Array shape a) where
202 |   cast (Scalar x) = x
203 |   cast [] = []
204 |   cast (x :: y) = cast @{toArray} x :: cast @{toArray} y
205 |
206 | namespace All
207 |   ||| An `All p xs` is an array (or scalar) of proofs about each element in `xs`.
208 |   |||
209 |   ||| For example, an `All IsSucc xs` proves that every element in `xs` is non-zero.
210 |   public export
211 |   data All : (0 p : a -> Type) -> Literal shape a -> Type where
212 |     Scalar : forall x . p x -> All p (Scalar x)
213 |     Nil  : All p []
214 |     (::) : All p x -> All p xs -> All p (x :: xs)
215 |
216 | namespace All2
217 |   ||| An `All2 p xs ys` is an array (or scalar) of pairwise proofs about elements in `xs` and `ys`.
218 |   |||
219 |   ||| For example, an `All2 LT xs ys` proves that each number in `xs` is less than the number in
220 |   ||| `ys` at the same position.
221 |   public export
222 |   data All2 : (p : a -> b -> Type) -> Literal shape a -> Literal shape b -> Type where
223 |     Scalar : forall a, b . p a b -> All2 p (Scalar a) (Scalar b)
224 |     Nil : All2 p [] []
225 |     (::) : All2 p a b -> All2 p as bs -> All2 p (a :: as) (b :: bs)
226 |