src/autograd/tensor_chain_LAWS.bend

bendtorch-core@0.1.0.0 · 175 lines · 17 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/autograd/gradient.bend as Gradient
5import bendtorch-core@0.1.0.0/src/autograd/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
27def Slots() -> Type:
28  Result<&2, &1, Backward.Error, List<&1, Gradient.Slot>>
29
30# Contributions sent so far, newest first.
31def Log() -> Type:
32  List<&1, Backward.Contribution>
33
34# The gathered slot of one id, and the contributions left for the others.
35def Taken() -> Type:
36  Result<&2, &1, Backward.Error, Gradient.Slot> & List<&1, Backward.Contribution>
37
38def 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
45def 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
53def 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.
64def 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
71def 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
78def 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.
85def 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.
96def 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
104def 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.
118def 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.
130def 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.
139law 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.
151def 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.
158law 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()}