PartitionedRegression.Coordinates

1module PartitionedRegression.Coordinates
2
3import Learning.ScaledGradientUpdate
4import Checkpoint.RegionManifest
5import Representation.Partition
6import Std.List
7import Std.Natural
8
9-- One prediction depends on both parameter owners. No participant can train
10-- its coordinate from its own prediction alone: its adjoint requires the
11-- residual of the sum. Device arithmetic is a separate admitted realization.
12def coordinatePrediction =
13  (lambda erased value : Type 0 .
14    (lambda unrestricted add : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
15      (lambda unrestricted multiply : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
16        (lambda unrestricted leftWeight : value .
17          (lambda unrestricted rightWeight : value .
18            (lambda unrestricted leftInput : value .
19              (lambda unrestricted rightInput : value .
20                (add (multiply leftWeight leftInput) (multiply rightWeight rightInput)))))))))
21
22-- Persistent vectors have the same logical coordinate domain. Their physical
23-- addresses and owners are supplied by placement; the epoch/update/cursor are
24-- the common capture prefix, not duplicated parameter fields.
25def coordinateCount : Nat =
26  2
27
28def coordinateWordBytes : Nat =
29  8
30
31def coordinateVectorBytes =
32  (naturalMultiply coordinateCount coordinateWordBytes)
33
34def coordinateRegions =
35  (constructor
36    StdList
37    StdListCons
38    (family CheckpointLogicalRegion)
39    (constructor
40      CheckpointLogicalRegion
41      CheckpointLogicalRegionOf
42      b"parameters"
43      b"f64-le"
44      coordinateVectorBytes)
45    (constructor
46      StdList
47      StdListCons
48      (family CheckpointLogicalRegion)
49      (constructor
50        CheckpointLogicalRegion
51        CheckpointLogicalRegionOf
52        b"scaled-gradient-accumulators"
53        b"f64-le"
54        coordinateVectorBytes)
55      (constructor StdList StdListEmpty (family CheckpointLogicalRegion))))
56
57def coordinateInterval =
58  (lambda unrestricted coordinate : Nat .
59    (constructor
60      LogicalInterval
61      LogicalIntervalOf
62      (naturalMultiply coordinate coordinateWordBytes)
63      (naturalMultiply (succ coordinate) coordinateWordBytes)))
64
65def coordinateIntervals =
66  (stdListMap
67    Nat
68    (family LogicalInterval)
69    coordinateInterval
70    (stdListNaturalsBelow coordinateCount))
71
72def coordinatePartitions =
73  (stdListMap
74    Nat
75    (family CheckpointRegionPartition)
76    (lambda unrestricted ordinal : Nat .
77      (constructor
78        CheckpointRegionPartition
79        CheckpointRegionPartitionOf
80        ordinal
81        coordinateIntervals))
82    (stdListNaturalsBelow (stdListLength (family CheckpointLogicalRegion) coordinateRegions)))