1	import Base
2	import ../shape/core.bend as Shape
3	import ../storage/main.bend as Storage
4	import ../tensor/core.bend as Core
5	import ../tensor/layout.bend as Layout
6	
7	# Gradient slots for reverse differentiation. A slot starts Unset. The first
8	# contribution whose dimensions agree with the slot's shape moves in unchanged;
9	# every later one is added in place, position by position in ascending flat
10	# order: total[i] = total[i] + contribution[i], one F32 add per position, never
11	# an overwrite. A contribution of any other shape is refused, and the slot and
12	# the contribution come back unchanged. Like PyTorch's .grad, grad reads an
13	# Unset slot as None; gradient instead materializes zeros of the slot's shape.
14	# reset returns a slot to Unset, like PyTorch's zero_grad(set_to_none=True).
15	
16	type Error is Data:
17	  TensorError{error: Core.Error}
18	  ShapeMismatch{}
19	
20	type Slot is Type:
21	  Unset{}
22	  Held{gradient: Core.Dynamic}
23	
24	type Outcome is Type:
25	  Accumulated{slot: Slot}
26	  Refused{slot: Slot, contribution: Core.Dynamic, error: Error}
27	
28	type State<-shape: Shape.Shape> is Type:
29	  State{contribution: Core.Dynamic, total: Core.Tensor<shape>}
30	
31	type Step<-shape: Shape.Shape> is Type:
32	  Continue{state: State<shape>}
33	  Stop{contribution: Core.Dynamic, total: Core.Tensor<shape>, error: Core.Error}
34	
35	def step_write(-shape: Shape.Shape, contribution: Core.Dynamic,
36	  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Step<shape>:
37	  (total, result) = pair
38	  match result:
39	    case Fail{error}:
40	      Stop{contribution, total, error}
41	    case Done{unit}:
42	      Continue{State{contribution, total}}
43	
44	def step_total(-shape: Shape.Shape, +index: U32, contribution: Core.Dynamic, value: F32,
45	  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Step<shape>:
46	  (total, result) = pair
47	  match result:
48	    case Fail{error}:
49	      Stop{contribution, total, error}
50	    case Done{+accumulated}:
51	      step_write(shape, contribution, Core.write_flat(shape, total, index, F32.add(accumulated, value)))
52	
53	def step_contribution(-shape: Shape.Shape, +index: U32, total: Core.Tensor<shape>,
54	  pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) -> Step<shape>:
55	  (contribution, result) = pair
56	  match result:
57	    case Fail{error}:
58	      Stop{contribution, total, error}
59	    case Done{value}:
60	      step_total(shape, index, contribution, value, Core.read_flat(shape, total, index))
61	
62	def step(-shape: Shape.Shape, +index: U32, state: State<shape>) -> Step<shape>:
63	  State{contribution, total} = state
64	  step_contribution(shape, index, total, Core.read_dynamic(contribution, index))
65	
66	def continue_step(+shape: Shape.Shape, rec: State<shape> -> Outcome, result: Step<shape>) -> Outcome:
67	  match result:
68	    case Stop{contribution, total, error}:
69	      Refused{Held{Core.Packed{shape, total}}, contribution, TensorError{error}}
70	    case Continue{state}:
71	      rec(state)
72	
73	def conversion_error(+shape: Shape.Shape, state: State<shape>) -> Outcome:
74	  State{contribution, total} = state
75	  Refused{Held{Core.Packed{shape, total}}, contribution, TensorError{Core.StorageError{Storage.InvalidIndex{}}}}
76	
77	def decided(+shape: Shape.Shape, +cursor: Nat, +index: U32, state: State<shape>, rec: State<shape> -> Outcome,
78	  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) -> Outcome:
79	  match decision:
80	    case Shape.BFalse{equal}:
81	      conversion_error(shape, state)
82	    case Shape.BTrue{equal}:
83	      continue_step(shape, rec, step(shape, index, state))
84	
85	def loop(fuel: Nat, +shape: Shape.Shape, +cursor: Nat, state: State<shape>) -> Outcome:
86	  match fuel:
87	    case 0n:
88	      State{contribution, total} = state
89	      Accumulated{Held{Core.Packed{shape, total}}}
90	    case 1n+remaining:
91	      +index = U32.from_nat(cursor)
92	      decided(shape, cursor, index, state, state => loop(remaining, shape, 1n+cursor, state),
93	        Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
94	
95	def held_checked(contribution: Core.Dynamic, +total_shape: Shape.Shape, total: Core.Tensor<total_shape>, same: Bool) -> Outcome:
96	  match same:
97	    case False{}:
98	      Refused{Held{Core.Packed{total_shape, total}}, contribution, ShapeMismatch{}}
99	    case True{}:
100	      loop(Shape.elements(total_shape), total_shape, 0n, State{contribution, total})
101	
102	def held(+shape: Shape.Shape, contribution: Core.Dynamic, gradient: Core.Dynamic) -> Outcome:
103	  Core.Packed{+total_shape, total} = gradient
104	  held_checked(contribution, total_shape, total, Layout.words_same(Shape.dimensions(total_shape), Shape.dimensions(shape)))
105	
106	def agreed(+shape: Shape.Shape, slot: Slot, contribution: Core.Dynamic) -> Outcome:
107	  match slot:
108	    case Unset{}:
109	      Accumulated{Held{contribution}}
110	    case Held{gradient}:
111	      held(shape, contribution, gradient)
112	
113	def checked(+shape: Shape.Shape, slot: Slot, +contribution_shape: Shape.Shape, tensor: Core.Tensor<contribution_shape>, same: Bool) -> Outcome:
114	  match same:
115	    case False{}:
116	      Refused{slot, Core.Packed{contribution_shape, tensor}, ShapeMismatch{}}
117	    case True{}:
118	      agreed(shape, slot, Core.Packed{contribution_shape, tensor})
119	
120	def accumulate(+shape: Shape.Shape, slot: Slot, contribution: Core.Dynamic) -> Outcome:
121	  Core.Packed{+contribution_shape, tensor} = contribution
122	  checked(shape, slot, contribution_shape, tensor,
123	    Layout.words_same(Shape.dimensions(contribution_shape), Shape.dimensions(shape)))
124	
125	# The gradient a slot holds, or None when it never received one or was reset.
126	def grad(slot: Slot) -> Maybe<&1, Core.Dynamic>:
127	  match slot:
128	    case Unset{}:
129	      None{}
130	    case Held{gradient}:
131	      Some{gradient}
132	
133	# Drops a held gradient: the next contribution moves in unchanged.
134	def reset(slot: Slot) -> Slot:
135	  Unset{}
136	
137	# The gradient a slot holds; zeros of shape when it never received one.
138	def gradient(+shape: Shape.Shape, slot: Slot) -> Result<&2, &1, Core.Error, Core.Dynamic>:
139	  match slot:
140	    case Unset{}:
141	      Core.construct_values(shape, Core.new(shape, 0.0))
142	    case Held{gradient}:
143	      Done{gradient}
144	
145	def error_code(error: Error) -> U32:
146	  match error:
147	    case TensorError{error}:
148	      Core.error_code(error)
149	    case ShapeMismatch{}:
150	      102