src/autograd/gradient_LAWS.bend

bendtorch-core@0.1.0.0 · 180 lines · 26 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality
5import bendtorch-core@0.1.0.0/src/autograd/gradient.bend as Gradient
6
7def add_word(total: Maybe<&2, F32>, value: Maybe<&2, F32>) -> Maybe<&2, F32>:
8  match total value:
9    case None{} _:
10      None{}
11    case Some{sum} None{}:
12      None{}
13    case Some{sum} Some{word}:
14      Some{F32.add(sum, word)}
15
16# The value the specification assigns to one position after a contribution is
17# accumulated into a slot: an Unset slot takes the contribution's word as it
18# is, and a Held slot adds the contribution's word to its own with one F32 add,
19# the held total first.
20def combined(slot: Gradient.Slot, contribution: Core.Dynamic, +index: U32) -> Maybe<&2, F32>:
21  match slot:
22    case Gradient.Unset{}:
23      Core.checked_observe_dynamic(contribution, index)
24    case Gradient.Held{gradient}:
25      add_word(Core.checked_observe_dynamic(gradient, index), Core.checked_observe_dynamic(contribution, index))
26
27def held_dimensions(slot: Gradient.Slot) -> Maybe<&2, List<&2, U32>>:
28  match slot:
29    case Gradient.Unset{}:
30      None{}
31    case Gradient.Held{gradient}:
32      Some{Totality.dimensions(gradient)}
33
34def held_observe(slot: Gradient.Slot, +index: U32) -> Maybe<&2, F32>:
35  match slot:
36    case Gradient.Unset{}:
37      None{}
38    case Gradient.Held{gradient}:
39      Core.checked_observe_dynamic(gradient, index)
40
41def Inside(+shape: Shape.Shape, +index: U32) -> Type:
42  {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}
43
44def Mapping(-slot: Gradient.Slot, -contribution: Core.Dynamic, +shape: Shape.Shape, outcome: Gradient.Outcome) -> Type:
45  match outcome:
46    case Gradient.Refused{after, returned, error}:
47      Unit
48    case Gradient.Accumulated{after}:
49      {held_dimensions(after) == Some{Shape.dimensions(shape)} : Maybe<&2, List<&2, U32>>} &
50        (@+index: U32 -> @inside: Inside(shape, index) ->
51          {held_observe(after, index) == combined(slot, contribution, index) : Maybe<&2, F32>})
52
53def SlotKept(before: Gradient.Slot, after: Gradient.Slot) -> Type:
54  match before:
55    case Gradient.Unset{}:
56      {held_dimensions(after) == None{} : Maybe<&2, List<&2, U32>>}
57    case Gradient.Held{gradient}:
58      {held_dimensions(after) == Some{Totality.dimensions(gradient)} : Maybe<&2, List<&2, U32>>} &
59        (@+index: U32 -> {held_observe(after, index) == Core.checked_observe_dynamic(gradient, index) : Maybe<&2, F32>})
60
61def mismatched(error: Gradient.Error) -> Bool:
62  match error:
63    case Gradient.TensorError{error}:
64      False{}
65    case Gradient.ShapeMismatch{}:
66      True{}
67
68def RefusalKeeps(-slot: Gradient.Slot, -contribution: Core.Dynamic, outcome: Gradient.Outcome) -> Type:
69  match outcome:
70    case Gradient.Accumulated{after}:
71      Unit
72    case Gradient.Refused{after, returned, error}:
73      {mismatched(error) == True{} : Bool} & (Totality.SameSource(contribution, returned) & SlotKept(slot, after))
74
75def Fits(+shape: Shape.Shape, slot: Gradient.Slot) -> Type:
76  match slot:
77    case Gradient.Unset{}:
78      Unit
79    case Gradient.Held{gradient}:
80      {Totality.dimensions(gradient) == Shape.dimensions(shape) : List<&2, U32>}
81
82def accumulated(outcome: Gradient.Outcome) -> Bool:
83  match outcome:
84    case Gradient.Accumulated{after}:
85      True{}
86    case Gradient.Refused{after, returned, error}:
87      False{}
88
89def Zeros(+shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Dynamic>) -> Type:
90  match result:
91    case Fail{error}:
92      Empty
93    case Done{gradient}:
94      {Totality.dimensions(gradient) == Shape.dimensions(shape) : List<&2, U32>} &
95        (@+index: U32 -> @inside: Inside(shape, index) ->
96          {Core.checked_observe_dynamic(gradient, index) == Some{0.0} : Maybe<&2, F32>})
97
98type MappingModel<-slot: Gradient.Slot, -contribution: Core.Dynamic, -shape: Shape.Shape, -outcome: Gradient.Outcome> is Data:
99  MappingModel{-mapping: Mapping(slot, contribution, shape, outcome)}
100
101type RefusalModel<-slot: Gradient.Slot, -contribution: Core.Dynamic, -outcome: Gradient.Outcome> is Data:
102  RefusalModel{-keeps: RefusalKeeps(slot, contribution, outcome)}
103
104type Success<-outcome: Gradient.Outcome> is Data:
105  Success{-successful: {accumulated(outcome) == True{} : Bool}}
106
107type ZerosModel<-shape: Shape.Shape, -result: Result<&2, &1, Core.Error, Core.Dynamic>> is Data:
108  ZerosModel{-zeros: Zeros(shape, result)}
109
110# A slot that never received a contribution reads as zeros of its shape.
111law gradient_unset_zeros:
112  for +shape: Shape.Shape
113  ZerosModel<shape, Gradient.gradient(shape, Gradient.Unset{})>
114
115# A slot that holds a gradient reads as that gradient.
116law gradient_held:
117  for +shape: Shape.Shape
118  for held: Core.Dynamic
119  {Gradient.gradient(shape, Gradient.Held{held}) == Done{held} : Result<&2, &1, Core.Error, Core.Dynamic>}
120
121# On success the slot holds a gradient with the slot's dimensions. At every
122# in-bounds position, the first contribution moves in unchanged and every
123# later one is added to the held total with one F32 add, total first.
124law accumulate_mapping:
125  for +shape: Shape.Shape
126  for slot: Gradient.Slot
127  for contribution: Core.Dynamic
128  MappingModel<slot, contribution, shape, Gradient.accumulate(shape, slot, contribution)>
129
130# Every refusal is a shape mismatch, decided before any write: the
131# contribution comes back unchanged, and so does the slot.
132law accumulate_refusal_keeps:
133  for +shape: Shape.Shape
134  for slot: Gradient.Slot
135  for contribution: Core.Dynamic
136  RefusalModel<slot, contribution, Gradient.accumulate(shape, slot, contribution)>
137
138# A contribution and a slot that both agree with the slot's shape always
139# accumulate.
140law accumulate_valid_success:
141  for +shape: Shape.Shape
142  for slot: Gradient.Slot
143  for contribution: Core.Dynamic
144  for fits: Fits(shape, slot)
145  for agrees: {Totality.dimensions(contribution) == Shape.dimensions(shape) : List<&2, U32>}
146  Success<Gradient.accumulate(shape, slot, contribution)>
147
148# Like PyTorch's .grad, a slot that never received a contribution has no
149# gradient: grad reads it as None, not as zeros.
150law grad_unset_none:
151  {Gradient.grad(Gradient.Unset{}) == None{} : Maybe<&1, Core.Dynamic>}
152
153# A slot that holds a gradient reads as that gradient.
154law grad_held_some:
155  for held: Core.Dynamic
156  {Gradient.grad(Gradient.Held{held}) == Some{held} : Maybe<&1, Core.Dynamic>}
157
158# reset returns every slot to Unset, like zero_grad(set_to_none=True).
159law reset_unset:
160  for slot: Gradient.Slot
161  {Gradient.reset(slot) == Gradient.Unset{} : Gradient.Slot}
162
163# After a reset, a contribution with the slot's dimensions moves in unchanged,
164# so a -0.0 word stays -0.0 instead of becoming +0.0 + -0.0.
165law reset_then_accumulate_moves_in:
166  for +shape: Shape.Shape
167  for slot: Gradient.Slot
168  for contribution: Core.Dynamic
169  for agrees: {Totality.dimensions(contribution) == Shape.dimensions(shape) : List<&2, U32>}
170  {Gradient.accumulate(shape, Gradient.reset(slot), contribution) == Gradient.Accumulated{Gradient.Held{contribution}} : Gradient.Outcome}
171
172# A contribution whose dimensions differ from the slot's is refused with
173# ShapeMismatch, and the slot and the contribution come back unchanged, even
174# when the element counts agree ([3, 2] or [6] for a [2, 3] slot).
175law accumulate_wrong_shape_refused:
176  for +shape: Shape.Shape
177  for slot: Gradient.Slot
178  for contribution: Core.Dynamic
179  for differ: {Totality.dimensions(contribution) != Shape.dimensions(shape) : List<&2, U32>}
180  {Gradient.accumulate(shape, slot, contribution) == Gradient.Refused{slot, contribution, Gradient.ShapeMismatch{}} : Gradient.Outcome}