type-safe name binding

spidr/spidr.ipkg

@@ -3,7 +3,11 @@ version = 0.0.6
 
 opts = "--total"
 sourcedir = "src"
-depends = elab-util
+
+depends =
+    elab-util
+  , linear
+
 modules =
   BayesianOptimization,
   BayesianOptimization.Acquisition,
@@ -32,6 +36,7 @@ modules =
   Data,
   Device,
   Distribution,
+  Linear,
   Literal,
 
   Model,

spidr/src/BayesianOptimization.idr

@@ -25,24 +25,26 @@ import public BayesianOptimization.Acquisition as BayesianOptimization
 import Data
 import Model
 import Tensor
+import Data.Linear.LVect
 
 ||| A `Stream`-like collection where each successive element can extend the set of `Tag`s.
 public export
 data TagStream : Type -> Type where
-  (::) : a -> Inf (Tag (TagStream a)) -> TagStream a
+  (::) : a -@ (Inf (Tag (TagStream a)) -> TagStream a)
 
 ||| Take `n` values from a `TagStream`, sequencing the `Tag` effects.
 public export
-take : (n : Nat) -> TagStream a -> Tag $ Vect n a
-take Z _ = pure Nil
+take : Copy a => (n : Nat) -> TagStream a -@ Tag $ LVect n a
+take Z (x :: _) = discarding x $ pure Nil
 take (S k) (x :: xs) = pure (x :: !(take k !xs))
 
 ||| Create an infinite stream of values from a generator function and a starting value.
 export covering
-iterate : (a -> Tag a) -> a -> Tag $ TagStream a
+iterate : Copy a => (a -@ Tag a) -> a -@ Tag $ TagStream a
 iterate f x = do
-  x' <- f x
-  pure (x' :: iterate f x')
+  x <- f x
+  MkBang x <- copy x
+  pure (x :: iterate f x)
 
 ||| Construct a single simple Bayesian optimization step.
 |||
@@ -51,13 +53,15 @@ iterate f x = do
 ||| @tactic The tactic, such as an optimized acquisition function, to find a new point from the
 |||   data and model
 export
-step : (objective : forall n . Tensor (n :: features) F64 -> Tag $ Tensor (n :: targets) F64) ->
+step : Copy model =>
+       (objective : forall n . Tensor (n :: features) F64 -@ Tag $ Tensor (n :: targets) F64) ->
        (probabilisticModel : ProbabilisticModel features targets marginal model) =>
-       (train : Dataset features targets -> model -> Tag $ model) ->
-       (tactic : ReaderT (DataModel {probabilisticModel} model) Tag (Tensor (1 :: features) F64)) ->
-       DataModel {probabilisticModel} model ->
-       Tag $ DataModel {probabilisticModel} model
+       (train : Dataset features targets -@ (model -@ Tag $ model)) ->
+       (tactic : DataModel {features, targets} model -@ Tag (Tensor (1 :: features) F64)) ->
+       DataModel {features, targets} model -@
+       Tag $ DataModel {features, targets} model
 step objective train tactic env = do
-  newPoint <- runReaderT env tactic
-  dataset <- tag $ concat env.dataset $ MkDataset newPoint !(objective newPoint)
+  MkBang env <- copy env
+  MkBang newPoint <- copy !(tactic env)
+  MkBang dataset <- copy $ concat env.dataset $ MkDataset newPoint !(objective newPoint)
   pure $ MkDataModel !(train dataset env.model) dataset

spidr/src/BayesianOptimization/Acquisition.idr

@@ -24,20 +24,26 @@ import Tensor
 import Data
 import Model
 
-%prefix_record_projections off
-
 ||| A `DataModel` packages data with a model over that data.
 public export
-record DataModel modelType {auto probabilisticModel : ProbabilisticModel f t marginal modelType} where
-  constructor MkDataModel
+data DataModel : {features, targets : Shape} -> Type -> Type where
+  MkDataModel : m -@ (Dataset features targets -@ DataModel {features, targets} m)
 
-  ||| A probabilistic model
-  model : modelType
+export
+Copy model => Copy (DataModel {features, targets} model) where
+  copy (MkDataModel m d) = do
+    MkBang m <- copy m
+    MkBang d <- copy d
+    pure $ MkBang $ MkDataModel m d
+  discard (MkDataModel m d) = discarding m $ discard d
 
-  ||| The data the model is trained on
-  dataset : Dataset f t
+export
+(.model) : DataModel m -@ m
+(MkDataModel m d).model = discarding d m
 
-%prefix_record_projections on
+export
+(.dataset) : Copy m => DataModel {features, targets} m -@ Dataset features targets
+(MkDataModel m d).dataset = discarding m d
 
 ||| An `Acquisition` function quantifies how useful it would be to query the objective at a given  
 ||| set of points, towards the goal of optimizing the objective.
@@ -47,7 +53,7 @@ record DataModel modelType {auto probabilisticModel : ProbabilisticModel f t mar
 ||| @features The shape of the feature domain.
 public export 0
 Acquisition : (0 batchSize : Nat) -> {auto 0 _ : GT batchSize 0} -> (0 features : Shape) -> Type
-Acquisition batchSize features = Tensor (batchSize :: features) F64 -> Tag $ Tensor [] F64
+Acquisition batchSize features = Tensor (batchSize :: features) F64 -@ Tag $ Tensor [] F64
 
 ||| Construct the acquisition function that estimates the absolute improvement in the best
 ||| observation if we were to evaluate the objective at a given point.
@@ -57,15 +63,15 @@ Acquisition batchSize features = Tensor (batchSize :: features) F64 -> Tag $ Ten
 export
 expectedImprovement :
   ProbabilisticModel features [1] Gaussian m =>
-  (model : m) ->
-  (best : Tensor [] F64) ->
+  (1 model : m) ->
+  (1 best : Tensor [] F64) ->
   Acquisition 1 features
 expectedImprovement model best at = do
-  best <- tag best
-  marginal <- tag =<< marginalise model at
+  MkBang best <- copy best
+  MkBang marginal <- copy !(marginalise model at)
   let best' = broadcast {to = [_, 1]} best
-  pdf <- tag =<< pdf marginal best'
-  cdf <- tag =<< cdf marginal best'
+  MkBang pdf <- copy !(pdf marginal best')
+  MkBang cdf <- copy !(cdf marginal best')
   let mean = squeeze !(mean {event = [1]} {dim = 1} marginal)
       variance = squeeze !(variance {event = [1]} marginal)
   pure $ (best - mean) * cdf + variance * pdf
@@ -74,12 +80,14 @@ expectedImprovement model best at = do
 ||| the observation value at each point.
 export
 expectedImprovementByModel :
+  Copy modelType =>
   ProbabilisticModel features [1] Gaussian modelType =>
-  ReaderT (DataModel modelType) Tag $ Acquisition 1 features
-expectedImprovementByModel = MkReaderT $ \env => do
-  marginal <- marginalise env.model env.dataset.features
-  best <- tag $ squeeze !(reduce @{Min} [0] !(mean {event = [1]} marginal))
-  pure $ expectedImprovement env.model best
+  DataModel {features, targets = [1]} modelType -@ Tag $ Acquisition 1 features
+expectedImprovementByModel (MkDataModel model (MkDataset features targets)) = do
+  MkBang model <- copy model
+  marginal <- marginalise model (discarding targets features)
+  MkBang best <- copy $ squeeze {to = []} !(reduce @{Min} [0] !(mean marginal))
+  pure $ expectedImprovement model best
 
 ||| Build an acquisition function that returns the probability that any given point will take a
 ||| value less than the specified `limit`.
@@ -88,9 +96,9 @@ probabilityOfFeasibility :
   (limit : Tensor [] F64) ->
   ClosedFormDistribution [1] dist =>
   ProbabilisticModel features [1] dist modelType =>
-  ReaderT (DataModel modelType) Tag $ Acquisition 1 features
-probabilityOfFeasibility limit =
-  asks $ \env, at => do cdf !(marginalise env.model at) (broadcast {to = [_, 1]} limit)
+  DataModel modelType -@ Acquisition 1 features
+probabilityOfFeasibility limit env at =
+  cdf !(marginalise env.model at) (broadcast {to = [_, 1]} limit)
 
 ||| Build an acquisition function that returns the negative of the lower confidence bound of the
 ||| probabilistic model. The variance contribution is weighted by a factor `beta`.
@@ -101,9 +109,9 @@ negativeLowerConfidenceBound :
   (beta : Double) ->
   {auto 0 betaNonNegative : beta >= 0 = True} ->
   ProbabilisticModel features [1] Gaussian modelType =>
-  ReaderT (DataModel modelType) Tag $ Acquisition 1 features
-negativeLowerConfidenceBound beta = asks $ \env, at => do
-  marginal <- tag =<< marginalise env.model at
+  DataModel modelType -@ Acquisition 1 features
+negativeLowerConfidenceBound beta env at = do
+  MkBang marginal <- copy !(marginalise env.model at)
   pure $ squeeze $
     !(mean {event = [1]} marginal) - fromDouble beta * !(variance {event = [1]} marginal)
 
@@ -117,4 +125,4 @@ export
 expectedConstrainedImprovement :
   (limit : Tensor [] F64) ->
   ProbabilisticModel features [1] Gaussian modelType =>
-  ReaderT (DataModel modelType) Tag (Acquisition 1 features -> Acquisition 1 features)
+  DataModel modelType -@ Tag (Acquisition 1 features -> Acquisition 1 features)

spidr/src/Compiler/Expr.idr

@@ -51,28 +51,6 @@ data Expr : Type
 export
 data Env = MkEnv Nat (List (Nat, Expr))
 
-export
-empty : Env
-empty = MkEnv 0 []
-
-export
-emptyFrom : Env -> Env
-emptyFrom (MkEnv n _) = MkEnv n []
-
-export
-updateCounterFrom : Env -> State Env ()
-updateCounterFrom (MkEnv n _) = do
-  MkEnv _ xs <- get
-  put $ MkEnv n xs
-
-export
-toList : Env -> List (Nat, Expr)
-toList (MkEnv _ env) = reverse env
-
-export
-counter : Env -> Nat
-counter (MkEnv c _) = c
-
 public export
 data Fn : Nat -> Type where
 
@@ -139,18 +117,32 @@ data Expr : Type where
   NormalFloatingPoint : (key, initialState : Expr) -> (shape : Shape) -> Expr
 
 export
-tag : Monad m => Expr -> StateT Env m Expr
-tag expr = do
-  MkEnv next env <- get
-  put $ MkEnv (S next) ((next, expr) :: env)
-  pure (Var next)
+empty : Env
+empty = MkEnv 0 []
+
+export
+emptyFrom : (1 _ : Env) -> Env
+emptyFrom (MkEnv n _) = MkEnv n []
+
+export
+updateCounterFrom : Env -> Env -> Env
+updateCounterFrom (MkEnv n _) (MkEnv _ xs) = MkEnv n xs
+
+export
+toList : Env -> List (Nat, Expr)
+toList (MkEnv _ env) = reverse env
+
+export
+counter : Env -> Nat
+counter (MkEnv c _) = c
+
+export
+copy : Expr -> (1 _ : Env) -> (Env, Expr)
+copy expr (MkEnv next env) = (MkEnv (S next) ((next, expr) :: env), Var next)
 
 export
-reserve : State Env Nat
-reserve = do
-  MkEnv next env <- get
-  put $ MkEnv (S next) env
-  pure next
+reserve : (1 _ : Env) -> (Env, Nat)
+reserve (MkEnv next env) = (MkEnv (S next) env, next)
 
 covering
 showExpr : Nat -> Expr -> String

spidr/src/Data.idr

@@ -18,32 +18,43 @@ module Data
 
 import Tensor
 
-%prefix_record_projections off
-
 ||| Observed pairs of data points from feature and target domains. Data sets such as this are
 ||| commonly used in supervised learning settings.
 |||
 ||| @features The shape of the feature domain.
 ||| @targets The shape of the target domain.
 public export
-record Dataset (0 featureShape, targetShape : Shape) where
-  constructor MkDataset
-  {s : Nat}
+data Dataset : (features, targets : Shape) -> Type where
+  MkDataset : {s : Nat} ->
+              Tensor (S s :: features) F64 -@ (
+              Tensor (S s :: targets) F64 -@
+              Dataset features targets)
 
-  ||| The feature data
-  features : Tensor (S s :: featureShape) F64
+public export
+size : Dataset f t -> Nat
+size (MkDataset {s} f t) = S s
 
-  ||| The target data
-  targets : Tensor (S s :: targetShape) F64
+||| The feature data
+export
+(.features) : (1 dataset : Dataset features targets) -> Tensor (size dataset :: features) F64
+(MkDataset fs ts).features = ts `seq` fs
 
-%prefix_record_projections on
+||| The target data
+export
+(.targets) : (1 dataset : Dataset features targets) -> Tensor (size dataset :: targets) F64
+(MkDataset fs ts).targets = fs `seq` ts
 
 ||| Concatenate two datasets along their leading axis.
 export
-concat : Dataset features targets -> Dataset features targets -> Dataset features targets
+concat : Dataset features targets -@ (Dataset features targets -@ Dataset features targets)
 concat (MkDataset {s = s} x y) (MkDataset {s = s'} x' y') =
   MkDataset {s = s + S s'} (concat 0 x x') (concat 0 y y')
 
 export
-Taggable (Dataset f t) where
-  tag (MkDataset f t) = [| MkDataset (tag f) (tag t) |]
+Copy (Dataset f t) where
+  copy (MkDataset f t) = pure $ zipWith MkDataset !(copy f) !(copy t)
+
+export
+Consumable (Dataset f t) where
+  consume (MkDataset f t) = f `seq` t `seq` ()
+