1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/elementwise_totality.bend as Totality
5	import ./gradient.bend as Gradient
6	
7	def 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.
20	def 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	
27	def 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	
34	def 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	
41	def Inside(+shape: Shape.Shape, +index: U32) -> Type:
42	  {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}
43	
44	def 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	
53	def 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	
61	def mismatched(error: Gradient.Error) -> Bool:
62	  match error:
63	    case Gradient.TensorError{error}:
64	      False{}
65	    case Gradient.ShapeMismatch{}:
66	      True{}
67	
68	def 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	
75	def 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	
82	def 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	
89	def 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	
98	type MappingModel<-slot: Gradient.Slot, -contribution: Core.Dynamic, -shape: Shape.Shape, -outcome: Gradient.Outcome> is Data:
99	  MappingModel{-mapping: Mapping(slot, contribution, shape, outcome)}
100	
101	type RefusalModel<-slot: Gradient.Slot, -contribution: Core.Dynamic, -outcome: Gradient.Outcome> is Data:
102	  RefusalModel{-keeps: RefusalKeeps(slot, contribution, outcome)}
103	
104	type Success<-outcome: Gradient.Outcome> is Data:
105	  Success{-successful: {accumulated(outcome) == True{} : Bool}}
106	
107	type 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.
111	law 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.
116	law 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.
124	law 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.
132	law 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.
140	law 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.
150	law 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.
154	law 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).
159	law 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.
165	law 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).
175	law 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}