1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ./gradient.bend as Gradient
5	import ./tensor_backward.bend as Backward
6	
7	# The gradient the whole sweep leaves, stated as the chain rule over the tape.
8	#
9	# The sweep pushes: it visits entries newest first, and an entry that holds a
10	# gradient sends its VJP's contributions into the working slots of its
11	# operands as it goes. The reference below pulls instead: it visits entries
12	# newest first, keeps every contribution sent so far in a log, newest first,
13	# and when it reaches the entry with id k it takes out of the log every
14	# contribution sent to k and accumulates them into k's starting slot, oldest
15	# first, with Gradient.accumulate. So the gradient of node k is its starting
16	# slot plus, in order, each contribution its consumers sent it: consumers
17	# newest first, and within one consumer in operand order. Each contribution is
18	# that consumer's Backward.vjp at the consumer's own gathered gradient, so the
19	# value of every contribution is the one its VJP value law states
20	# (tensor_vjp_LAWS.bend). The order of the float sum is part of the statement.
21	#
22	# The reference shares with the sweep only Backward.vjp and
23	# Gradient.accumulate, each stated by its own laws, and the list helpers
24	# Backward.prepend and Backward.unsets. It has no slot list that contributions
25	# are pushed into, no delivery and no at_head bookkeeping.
26	
27	def Slots() -> Type:
28	  Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>
29	
30	# Contributions sent so far, newest first.
31	def Log() -> Type:
32	  List<&1, Backward.Contribution>
33	
34	# The gathered slot of one id, and the contributions left for the others.
35	def Taken() -> Type:
36	  Result<&2, &1, Backward.Error, Gradient.Slot> & List<&1, Backward.Contribution>
37	
38	def settled(outcome: Gradient.Outcome) -> Result<&2, &1, Backward.Error, Gradient.Slot>:
39	  match outcome:
40	    case Gradient.Accumulated{slot}:
41	      Done{slot}
42	    case Gradient.Refused{slot, contribution, error}:
43	      Fail{Backward.GradientError{error}}
44	
45	def poured(+shape: Shape.Shape, held: Result<&2, &1, Backward.Error, Gradient.Slot>, gradient: Core.Dynamic) ->
46	  Result<&2, &1, Backward.Error, Gradient.Slot>:
47	  match held:
48	    case Fail{error}:
49	      Fail{error}
50	    case Done{slot}:
51	      settled(Gradient.accumulate(shape, slot, gradient))
52	
53	def sorted(+target: Nat, +shape: Shape.Shape, gradient: Core.Dynamic, mine: Bool, pair: Taken()) -> Taken():
54	  match mine:
55	    case True{}:
56	      (held, rest) = pair
57	      (poured(shape, held, gradient), rest)
58	    case False{}:
59	      (held, rest) = pair
60	      (held, Backward.Contribution{target, shape, gradient} <> rest)
61	
62	# Takes out of a newest-first log every contribution to id and accumulates
63	# them into slot, oldest first; the others stay in the log, in order.
64	def take(log: Log(), +id: Nat, slot: Gradient.Slot) -> Taken():
65	  match log:
66	    case Nil{}:
67	      (Done{slot}, Nil{})
68	    case Con{Backward.Contribution{+target, +shape, gradient}, rest}:
69	      sorted(target, shape, gradient, Nat.is_eq(id, target), take(rest, id, slot))
70	
71	def joined(held: Result<&2, &1, Backward.Error, Gradient.Slot>, rest: Slots()) -> Slots():
72	  match held:
73	    case Fail{error}:
74	      Fail{error}
75	    case Done{slot}:
76	      Backward.prepend(Gradient.Slot, slot, rest)
77	
78	def gather_step(pair: Taken(), rec: Log() -> Slots()) -> Slots():
79	  (held, rest) = pair
80	  joined(held, rec(rest))
81	
82	# Each slot of a newest-first list of count slots, the one with k older slots
83	# for id k, with what the log sends to it. A list longer than count keeps its
84	# slots past count, as the sweep does.
85	def gather(slots: List<&1, Gradient.Slot>, +count: Nat, log: Log()) -> Slots():
86	  match slots count:
87	    case Nil{} count:
88	      Done{Nil{}}
89	    case Con{slot, rest} 0n:
90	      Done{slot <> rest}
91	    case Con{slot, rest} 1n+below:
92	      gather_step(take(log, below, slot), more => gather(rest, below, more))
93	
94	# An entry that holds g after gathering: its VJP's contributions join the log
95	# in operand order, so the first operand's arrives first.
96	def fired(pass: Backward.Pass, log: Log(), rec: Log() -> Slots()) -> Slots():
97	  Backward.Pass{saved, g, contributions} = pass
98	  match contributions:
99	    case Fail{error}:
100	      Fail{error}
101	    case Done{items}:
102	      Backward.prepend(Gradient.Slot, Gradient.Held{g}, rec(List.reverse.go(&1, Backward.Contribution, items, log)))
103	
104	def reached(node: Backward.Node, saved: Backward.Saved, pair: Taken(), rec: Log() -> Slots()) -> Slots():
105	  (held, log) = pair
106	  match held:
107	    case Fail{error}:
108	      Fail{error}
109	    case Done{Gradient.Unset{}}:
110	      Backward.prepend(Gradient.Slot, Gradient.Unset{}, rec(log))
111	    case Done{Gradient.Held{g}}:
112	      fired(Backward.vjp(node, saved, g), log, rec)
113	
114	# The working gradients by the chain rule, pulled, for entries newest first
115	# that start with slots and receive the contributions in log: the entry with
116	# id k gathers every contribution sent to k into its starting slot; if it then
117	# holds g, the contributions of vjp(node, saved, g) join the log.
118	def chained(entries: List<&1, Backward.Entry>, slots: List<&1, Gradient.Slot>, +count: Nat, log: Log()) -> Slots():
119	  match entries slots count:
120	    case Nil{} slots count:
121	      gather(slots, count, log)
122	    case Con{entry, rest} Nil{} count:
123	      Done{Nil{}}
124	    case Con{entry, rest} Con{slot, more} 0n:
125	      Done{Nil{}}
126	    case Con{Backward.Entry{node, shape, value, saved, grad}, rest} Con{slot, more} 1n+below:
127	      reached(node, saved, take(log, below, slot), next => chained(rest, more, below, next))
128	
129	# The working gradients of a sweep, without its entries.
130	def result(pair: Backward.Swept()) -> Slots():
131	  (entries, result) = pair
132	  result
133	
134	# The chain rule, for any recorded tape and any starting working gradients:
135	# when the sweep succeeds, the working gradient it leaves at every entry is
136	# the chained one, each slot the accumulation, in the order above, of the VJP
137	# values its consumers sent it at their own chained gradients. A refusal is
138	# stated by the structure laws in tensor_backward_LAWS.bend.
139	law tensor_sweep_chain_rule:
140	  for entries: List<&1, Backward.Entry>
141	  for gradients: List<&1, Gradient.Slot>
142	  for +count: Nat
143	  for -swept: List<&1, Gradient.Slot>
144	  for succeeded: {result(Backward.sweep(entries, Done{gradients}, count)) == Done{swept} : Slots()}
145	  {chained(entries, gradients, count, Nil{}) == Done{swept} : Slots()}
146	
147	# The starting working gradients of a backward from the entry with older
148	# entries below it and newer entries above it: no gradient for the newer
149	# entries, ones of the output's shape at the output, as PyTorch's
150	# backward(torch.ones_like(output)), and no gradient below.
151	def seeds(+newer: Nat, +shape: Shape.Shape, ones: Core.Tensor<shape>, +older: Nat) -> List<&1, Gradient.Slot>:
152	  List.append(&1, Gradient.Slot, Backward.unsets(newer), Gradient.Held{Core.Packed{shape, ones}} <> Backward.unsets(older))
153	
154	# The chain rule for the whole backward: when the pass of a backward from the
155	# output with older entries below it succeeds, the working gradients it hands
156	# to the merge (tensor_backward_success, tensor_merge_accumulates) are the
157	# chained ones from the seeds, for any tape, newer entries included.
158	law tensor_backward_chain_rule:
159	  for newer: List<&1, Backward.Entry>
160	  for +node: Backward.Node
161	  for +shape: Shape.Shape
162	  for value: Core.Dynamic
163	  for saved: Backward.Saved
164	  for slot: Gradient.Slot
165	  for rest: List<&1, Backward.Entry>
166	  for +older: Nat
167	  for ones: Core.Tensor<shape>
168	  for made: {Core.new(shape, 1.0) == Done{ones} : Result<&2, &1, Core.Error, Core.Tensor<shape>>}
169	  for -gradients: List<&1, Gradient.Slot>
170	  for succeeded: {result(Backward.pass(List.append(&1, Backward.Entry, newer, Backward.Entry{node, shape, value, saved, slot} <> rest),
171	      1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), older,
172	      Backward.at_head(1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), older))) == Done{gradients} : Slots()}
173	  {chained(List.append(&1, Backward.Entry, newer, Backward.Entry{node, shape, value, saved, slot} <> rest),
174	      seeds(List.length(&1, Backward.Entry, newer), shape, ones, older),
175	      1n+Nat.add(older, List.length(&1, Backward.Entry, newer)), Nil{}) == Done{gradients} : Slots()}