1 import Base 2 import ./backward.bend as Backward 3 4 # Reachability, stated on flags: a node is reached when it is the selected 5 # output, or an operand of a reached node newer than it. Nodes newer than the 6 # output are never reached. 7 8 def marked(flag: Bool, +rest: List<&2, Bool>, rec: Unit -> List<&2, Bool>, same: Bool) -> List<&2, Bool>: 9 match same: 10 case True{}: 11 True{} <> rest 12 case False{}: 13 flag <> rec(Unit{}) 14 15 def mark(flags: List<&2, Bool>, +count: Nat, +target: Nat) -> List<&2, Bool>: 16 match flags count: 17 case Nil{} count: 18 Nil{} 19 case Con{flag, rest} 0n: 20 flag <> rest 21 case Con{flag, +rest} 1n+below: 22 marked(flag, rest, unit => mark(rest, below, target), Nat.is_eq(below, target)) 23 24 def targets(node: Backward.Node) -> List<&2, Nat>: 25 match node: 26 case Backward.Leaf{}: 27 [] 28 case Backward.Add{a, b}: 29 [a, b] 30 case Backward.Sub{a, b}: 31 [a, b] 32 case Backward.Mul{a, b, left, right}: 33 [a, b] 34 case Backward.Neg{a}: 35 [a] 36 37 def mark_all(ids: List<&2, Nat>, +count: Nat, flags: List<&2, Bool>) -> List<&2, Bool>: 38 match ids: 39 case Nil{}: 40 flags 41 case Con{+id, rest}: 42 mark_all(rest, count, mark(flags, count, id)) 43 44 def spread(node: Backward.Node, flag: Bool, +count: Nat, flags: List<&2, Bool>) -> List<&2, Bool>: 45 match flag: 46 case False{}: 47 flags 48 case True{}: 49 mark_all(targets(node), count, flags) 50 51 def reach(nodes: List<&2, Backward.Node>, flags: List<&2, Bool>, +count: Nat) -> List<&2, Bool>: 52 match nodes flags count: 53 case Nil{} flags count: 54 flags 55 case Con{node, rest} Nil{} count: 56 Nil{} 57 case Con{node, rest} Con{flag, pending} 0n: 58 Nil{} 59 case Con{node, rest} Con{+flag, pending} 1n+below: 60 flag <> reach(rest, spread(node, flag, below, pending), below) 61 62 def falses(count: Nat) -> List<&2, Bool>: 63 match count: 64 case 0n: 65 Nil{} 66 case 1n+rest: 67 False{} <> falses(rest) 68 69 def reached_chosen(node: Backward.Node, +rest: List<&2, Backward.Node>, +below: Nat, rec: Unit -> List<&2, Bool>, selected: Bool) -> 70 List<&2, Bool>: 71 match selected: 72 case True{}: 73 reach(node <> rest, True{} <> falses(below), 1n+below) 74 case False{}: 75 False{} <> rec(Unit{}) 76 77 def reached(nodes: List<&2, Backward.Node>, +count: Nat, +output: Nat) -> List<&2, Bool>: 78 match nodes count: 79 case Nil{} count: 80 Nil{} 81 case Con{node, rest} 0n: 82 Nil{} 83 case Con{node, +rest} 1n+below: 84 reached_chosen(node, rest, below, unit => reached(rest, below, output), Nat.is_eq(below, output)) 85 86 # Operand bounds, stated without the engine's validator: a node may name only 87 # nodes older than itself, so each operand id must be below the number of 88 # older nodes. An operand that is the node itself, a newer node or a missing 89 # node fails this bound. 90 def bounded(ids: List<&2, Nat>, +bound: Nat) -> Bool: 91 match ids: 92 case Nil{}: 93 True{} 94 case Con{+id, rest}: 95 Bool.and(Nat.is_lt(id, bound), bounded(rest, bound)) 96 97 def well_formed(nodes: List<&2, Backward.Node>) -> Bool: 98 match nodes: 99 case Nil{}: 100 True{} 101 case Con{node, +rest}: 102 Bool.and(bounded(targets(node), List.length(&2, Backward.Node, rest)), well_formed(rest)) 103 104 def entry_node(entry: Backward.Entry) -> Backward.Node: 105 Backward.Entry{node, value, slot} = entry 106 node 107 108 def somes(gradients: List<&2, Maybe<&2, F32>>) -> List<&2, Bool>: 109 match gradients: 110 case Nil{}: 111 Nil{} 112 case Con{gradient, rest}: 113 Maybe.is_some(&2, F32, gradient) <> somes(rest) 114 115 def nodes_of(tape: Backward.Tape) -> List<&2, Backward.Node>: 116 Backward.Tape{entries} = tape 117 Backward.nodes(entries) 118 119 def slots(entries: List<&2, Backward.Entry>) -> List<&2, Maybe<&2, F32>>: 120 match entries: 121 case Nil{}: 122 Nil{} 123 case Con{Backward.Entry{node, value, slot}, rest}: 124 slot <> slots(rest) 125 126 def tape_slots(tape: Backward.Tape) -> List<&2, Maybe<&2, F32>>: 127 Backward.Tape{entries} = tape 128 slots(entries) 129 130 def slots_after(result: Result<&2, &2, Backward.Error, Backward.Tape>) -> Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>: 131 match result: 132 case Fail{error}: 133 Fail{error} 134 case Done{tape}: 135 Done{tape_slots(tape)} 136 137 def grad_after(result: Result<&2, &2, Backward.Error, Backward.Tape>, +id: Nat) -> Result<&2, &2, Backward.Error, Maybe<&2, F32>>: 138 match result: 139 case Fail{error}: 140 Fail{error} 141 case Done{tape}: 142 Backward.grad(tape, id) 143 144 def again(result: Result<&2, &2, Backward.Error, Backward.Tape>, +output: Nat) -> Result<&2, &2, Backward.Error, Backward.Tape>: 145 match result: 146 case Fail{error}: 147 Fail{error} 148 case Done{tape}: 149 Backward.backward(tape, output) 150 151 def observe(+tape: Backward.Tape, +output: Nat) -> 152 Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>: 153 (Backward.value(tape, output), Backward.gradients(tape, output)) 154 155 def tape_of(pair: Backward.Tape & Nat) -> Backward.Tape: 156 (tape, id) = pair 157 tape 158 159 def recorded(pair: Backward.Tape & Result<&2, &2, Backward.Error, Nat>) -> Backward.Tape: 160 (tape, result) = pair 161 tape 162 163 def program_binary(kind: Backward.Binary, x: F32, y: F32) -> Backward.Tape: 164 recorded(Backward.binary(tape_of(Backward.leaf(tape_of(Backward.leaf(Backward.empty(), x)), y)), kind, 0n, 1n)) 165 166 def program_neg(x: F32) -> Backward.Tape: 167 recorded(Backward.neg(tape_of(Backward.leaf(Backward.empty(), x)), 0n)) 168 169 # x = 0, y = 1, z = 2, x * y = 3, z + z = 4. Node 4 is newer than the output 170 # 3 and uses z twice; z is not used by 3. 171 def program_unused(x: F32, y: F32, z: F32) -> Backward.Tape: 172 recorded(Backward.add(recorded(Backward.mul( 173 tape_of(Backward.leaf(tape_of(Backward.leaf(tape_of(Backward.leaf(Backward.empty(), x)), y)), z)), 0n, 1n)), 2n, 2n)) 174 175 # x = 0, a = 1, b = 2, c = 3, p = x * a = 4, q = x * b = 5, r = x * c = 6, 176 # s = p + q = 7, t = s + r = 8. x has three consumers. 177 def program_fan_out(x: F32, a: F32, b: F32, c: F32) -> Backward.Tape: 178 recorded(Backward.add(recorded(Backward.add(recorded(Backward.mul(recorded(Backward.mul(recorded(Backward.mul( 179 tape_of(Backward.leaf(tape_of(Backward.leaf(tape_of(Backward.leaf(tape_of(Backward.leaf(Backward.empty(), x)), a)), b)), c)), 180 0n, 1n)), 0n, 2n)), 0n, 3n)), 4n, 5n)), 7n, 6n)) 181 182 # Vector-Jacobian products: the contribution each operand receives when the 183 # node's gradient is g, in operand order. 184 185 law vjp_leaf: 186 for +g: F32 187 {Backward.vjp(Backward.Leaf{}, g) == [] : List<&2, Backward.Contribution>} 188 189 law vjp_add: 190 for +a: Nat 191 for +b: Nat 192 for +g: F32 193 {Backward.vjp(Backward.Add{a, b}, g) == [Backward.Contribution{a, g}, Backward.Contribution{b, g}] : List<&2, Backward.Contribution>} 194 195 law vjp_sub: 196 for +a: Nat 197 for +b: Nat 198 for +g: F32 199 {Backward.vjp(Backward.Sub{a, b}, g) == [Backward.Contribution{a, g}, Backward.Contribution{b, F32.neg(g)}] : List<&2, Backward.Contribution>} 200 201 # d(a*b)/da = b and d(a*b)/db = a, with the operand values saved at record 202 # time: the first operand receives g * right, the second g * left. 203 law vjp_mul: 204 for +a: Nat 205 for +b: Nat 206 for +left: F32 207 for +right: F32 208 for +g: F32 209 {Backward.vjp(Backward.Mul{a, b, left, right}, g) == 210 [Backward.Contribution{a, F32.mul(g, right)}, Backward.Contribution{b, F32.mul(g, left)}] : List<&2, Backward.Contribution>} 211 212 law vjp_neg: 213 for +a: Nat 214 for +g: F32 215 {Backward.vjp(Backward.Neg{a}, g) == [Backward.Contribution{a, F32.neg(g)}] : List<&2, Backward.Contribution>} 216 217 # Recording through the public API and running backward gives each op's 218 # forward value and its VJP at the leaves (gradients newest first). 219 220 law backward_add: 221 for +x: F32 222 for +y: F32 223 {observe(program_binary(Backward.Plus{}, x, y), 2n) == 224 (Done{F32.add(x, y)}, Done{[Some{1.0}, Some{1.0}, Some{1.0}]}) : 225 Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>} 226 227 law backward_sub: 228 for +x: F32 229 for +y: F32 230 {observe(program_binary(Backward.Minus{}, x, y), 2n) == 231 (Done{F32.sub(x, y)}, Done{[Some{1.0}, Some{F32.neg(1.0)}, Some{1.0}]}) : 232 Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>} 233 234 law backward_mul: 235 for +x: F32 236 for +y: F32 237 {observe(program_binary(Backward.Times{}, x, y), 2n) == 238 (Done{F32.mul(x, y)}, Done{[Some{1.0}, Some{F32.mul(1.0, x)}, Some{F32.mul(1.0, y)}]}) : 239 Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>} 240 241 law backward_neg: 242 for +x: F32 243 {observe(program_neg(x), 1n) == 244 (Done{F32.neg(x)}, Done{[Some{1.0}, Some{F32.neg(1.0)}]}) : 245 Result<&2, &2, Backward.Error, F32> & Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>} 246 247 # The seed is the selected output (C1). For any graph, the nodes newer than 248 # the output get None, the output gets exactly 1.0, and the older nodes get 249 # what the sweep sends from the output's VJP, whatever the newer nodes are. 250 law seed_selected_output: 251 for +newer: List<&2, Backward.Node> 252 for +node: Backward.Node 253 for +rest: List<&2, Backward.Node> 254 {Backward.pass(List.append(&2, Backward.Node, newer, node <> rest), 255 List.length(&2, Backward.Node, List.append(&2, Backward.Node, newer, node <> rest)), List.length(&2, Backward.Node, rest)) == 256 List.append(&2, Maybe<&2, F32>, Backward.nones(List.length(&2, Backward.Node, newer)), 257 Some{1.0} <> Backward.sweep(rest, 258 Backward.contribute(node, Some{1.0}, List.length(&2, Backward.Node, rest), Backward.nones(List.length(&2, Backward.Node, rest))), 259 List.length(&2, Backward.Node, rest))) : List<&2, Maybe<&2, F32>>} 260 261 # Only the reachable graph: a node has a gradient exactly when the selected 262 # output reaches it. Every other node, newer or unused, reads None. 263 law reached_exactly: 264 for +graph: List<&2, Backward.Node> 265 for +output: Nat 266 {somes(Backward.pass(graph, List.length(&2, Backward.Node, graph), output)) == 267 reached(graph, List.length(&2, Backward.Node, graph), output) : List<&2, Bool>} 268 269 # With x, y and z recorded and z only used by a node newer than the output 270 # x * y, backward of x * y seeds the output, leaves z and the newer node None, 271 # and gives x and y their VJP. 272 law unused_leaf_none: 273 for +x: F32 274 for +y: F32 275 for +z: F32 276 {slots_after(Backward.backward(program_unused(x, y, z), 3n)) == 277 Done{[None{}, Some{1.0}, None{}, Some{F32.mul(1.0, x)}, Some{F32.mul(1.0, y)}]} : 278 Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>} 279 280 # Fan-out: a node's first contribution moves in, and each later one is added 281 # to the held total with one F32 add, total first. 282 law accumulate_moves_in: 283 for contribution: Maybe<&2, F32> 284 {Backward.accumulate(None{}, contribution) == contribution : Maybe<&2, F32>} 285 286 law accumulate_adds_total_first: 287 for +held: F32 288 for +value: F32 289 {Backward.accumulate(Some{held}, Some{value}) == Some{F32.add(held, value)} : Maybe<&2, F32>} 290 291 law accumulate_keeps_without_contribution: 292 for +held: F32 293 {Backward.accumulate(Some{held}, None{}) == Some{held} : Maybe<&2, F32>} 294 295 # A contribution to node target reaches exactly the slot of that node and no 296 # other, for any list of slots. 297 law bump_target: 298 for +newer: List<&2, Maybe<&2, F32>> 299 for +slot: Maybe<&2, F32> 300 for +rest: List<&2, Maybe<&2, F32>> 301 for +value: F32 302 {Backward.bump(List.append(&2, Maybe<&2, F32>, newer, slot <> rest), 303 List.length(&2, Maybe<&2, F32>, List.append(&2, Maybe<&2, F32>, newer, slot <> rest)), List.length(&2, Maybe<&2, F32>, rest), value) == 304 List.append(&2, Maybe<&2, F32>, newer, Backward.accumulate(slot, Some{value}) <> rest) : List<&2, Maybe<&2, F32>>} 305 306 # The contributions of x's three consumers arrive newest consumer first, and 307 # each is added to the held total. 308 law fan_out_arrival_order: 309 for +x: F32 310 for +a: F32 311 for +b: F32 312 for +c: F32 313 {grad_after(Backward.backward(program_fan_out(x, a, b, c), 8n), 0n) == 314 Done{Some{F32.add(F32.add(F32.mul(1.0, c), F32.mul(1.0, b)), F32.mul(1.0, a))}} : 315 Result<&2, &2, Backward.Error, Maybe<&2, F32>>} 316 317 # Ids (C2): id k reads exactly the element with k older elements, and an id 318 # the list does not hold reads None, never another element. 319 law lookup_exact: 320 for -A: Data 321 for +newer: List<&2, A> 322 for +item: A 323 for +rest: List<&2, A> 324 {Backward.lookup(A, List.append(&2, A, newer, item <> rest), List.length(&2, A, rest)) == Some{item} : Maybe<&2, A>} 325 326 law lookup_unknown: 327 for -A: Data 328 for +items: List<&2, A> 329 for +id: Nat 330 for outside: {Nat.is_lt(id, List.length(&2, A, items)) == False{} : Bool} 331 {Backward.lookup(A, items, id) == None{} : Maybe<&2, A>} 332 333 law bump_unknown: 334 for +gradients: List<&2, Maybe<&2, F32>> 335 for +target: Nat 336 for +value: F32 337 for +outside: {Nat.is_lt(target, List.length(&2, Maybe<&2, F32>, gradients)) == False{} : Bool} 338 {Backward.bump(gradients, List.length(&2, Maybe<&2, F32>, gradients), target, value) == gradients : List<&2, Maybe<&2, F32>>} 339 340 law value_unknown: 341 for +entries: List<&2, Backward.Entry> 342 for +id: Nat 343 for outside: {Nat.is_lt(id, List.length(&2, Backward.Entry, entries)) == False{} : Bool} 344 {Backward.value(Backward.Tape{entries}, id) == Fail{Backward.UnknownNode{}} : Result<&2, &2, Backward.Error, F32>} 345 346 law grad_unknown: 347 for +entries: List<&2, Backward.Entry> 348 for +id: Nat 349 for outside: {Nat.is_lt(id, List.length(&2, Backward.Entry, entries)) == False{} : Bool} 350 {Backward.grad(Backward.Tape{entries}, id) == Fail{Backward.UnknownNode{}} : Result<&2, &2, Backward.Error, Maybe<&2, F32>>} 351 352 # Recording with an operand id the tape does not hold is refused with 353 # UnknownNode, and the tape comes back unchanged. 354 law record_unknown_left: 355 for +entries: List<&2, Backward.Entry> 356 for kind: Backward.Binary 357 for +a: Nat 358 for +b: Nat 359 for outside: {Nat.is_lt(a, List.length(&2, Backward.Entry, entries)) == False{} : Bool} 360 {Backward.binary(Backward.Tape{entries}, kind, a, b) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) : 361 Backward.Tape & Result<&2, &2, Backward.Error, Nat>} 362 363 law record_unknown_right: 364 for +entries: List<&2, Backward.Entry> 365 for kind: Backward.Binary 366 for +a: Nat 367 for +b: Nat 368 for outside: {Nat.is_lt(b, List.length(&2, Backward.Entry, entries)) == False{} : Bool} 369 {Backward.binary(Backward.Tape{entries}, kind, a, b) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) : 370 Backward.Tape & Result<&2, &2, Backward.Error, Nat>} 371 372 law neg_unknown: 373 for +entries: List<&2, Backward.Entry> 374 for +a: Nat 375 for outside: {Nat.is_lt(a, List.length(&2, Backward.Entry, entries)) == False{} : Bool} 376 {Backward.neg(Backward.Tape{entries}, a) == (Backward.Tape{entries}, Fail{Backward.UnknownNode{}}) : 377 Backward.Tape & Result<&2, &2, Backward.Error, Nat>} 378 379 # backward of an output the tape does not hold is refused with UnknownNode. 380 law backward_unknown: 381 for +entries: List<&2, Backward.Entry> 382 for +output: Nat 383 for outside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == False{} : Bool} 384 {Backward.backward(Backward.Tape{entries}, output) == Fail{Backward.UnknownNode{}} : Result<&2, &2, Backward.Error, Backward.Tape>} 385 386 # A tape with a node that refers to itself, to a newer node or to a missing 387 # node is refused with DanglingOperand. 388 law backward_dangling: 389 for +entries: List<&2, Backward.Entry> 390 for +output: Nat 391 for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == True{} : Bool} 392 for unsound: {Backward.wired(Backward.nodes(entries), List.length(&2, Backward.Entry, entries)) == False{} : Bool} 393 {Backward.backward(Backward.Tape{entries}, output) == Fail{Backward.DanglingOperand{}} : Result<&2, &2, Backward.Error, Backward.Tape>} 394 395 # The engine's wiring check is exactly the operand-bound specification above. 396 law wired_matches_operand_bounds: 397 for +graph: List<&2, Backward.Node> 398 {Backward.wired(graph, List.length(&2, Backward.Node, graph)) == well_formed(graph) : Bool} 399 400 # An entry with an operand that is itself, a newer entry or a missing entry 401 # makes backward of any output on the tape fail with DanglingOperand. 402 law backward_rejects_unbound_operand: 403 for +newer: List<&2, Backward.Entry> 404 for +entry: Backward.Entry 405 for +rest: List<&2, Backward.Entry> 406 for +output: Nat 407 for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, List.append(&2, Backward.Entry, newer, entry <> rest))) == True{} : Bool} 408 for unbound: {bounded(targets(entry_node(entry)), List.length(&2, Backward.Entry, rest)) == False{} : Bool} 409 {Backward.backward(Backward.Tape{List.append(&2, Backward.Entry, newer, entry <> rest)}, output) == Fail{Backward.DanglingOperand{}} : 410 Result<&2, &2, Backward.Error, Backward.Tape>} 411 412 # A tape whose every operand is bound is accepted for every output it holds. 413 law backward_accepts_bound_operands: 414 for +entries: List<&2, Backward.Entry> 415 for +output: Nat 416 for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == True{} : Bool} 417 for bound: {well_formed(Backward.nodes(entries)) == True{} : Bool} 418 {Backward.backward(Backward.Tape{entries}, output) == 419 Done{Backward.Tape{Backward.merge(entries, Backward.pass(Backward.nodes(entries), List.length(&2, Backward.Entry, entries), output))}} : 420 Result<&2, &2, Backward.Error, Backward.Tape>} 421 422 # Otherwise backward accumulates the gradients of pass into the slots. 423 law backward_success: 424 for +entries: List<&2, Backward.Entry> 425 for +output: Nat 426 for inside: {Nat.is_lt(output, List.length(&2, Backward.Entry, entries)) == True{} : Bool} 427 for sound: {Backward.wired(Backward.nodes(entries), List.length(&2, Backward.Entry, entries)) == True{} : Bool} 428 {Backward.backward(Backward.Tape{entries}, output) == 429 Done{Backward.Tape{Backward.merge(entries, Backward.pass(Backward.nodes(entries), List.length(&2, Backward.Entry, entries), output))}} : 430 Result<&2, &2, Backward.Error, Backward.Tape>} 431 432 # Repeated backward, like PyTorch's backward(retain_graph=True): the second 433 # backward computes the same gradients and accumulates them again. 434 law backward_twice_accumulates: 435 for +entries: List<&2, Backward.Entry> 436 for +output: Nat 437 for +found: List<&2, Maybe<&2, F32>> 438 for +computed: {Backward.gradients(Backward.Tape{entries}, output) == Done{found} : Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>} 439 {again(Backward.backward(Backward.Tape{entries}, output), output) == 440 Done{Backward.Tape{Backward.merge(Backward.merge(entries, found), found)}} : Result<&2, &2, Backward.Error, Backward.Tape>} 441 442 law backward_twice_mul: 443 for +x: F32 444 for +y: F32 445 {slots_after(again(Backward.backward(program_binary(Backward.Times{}, x, y), 2n), 2n)) == 446 Done{[Some{F32.add(1.0, 1.0)}, Some{F32.add(F32.mul(1.0, x), F32.mul(1.0, x))}, Some{F32.add(F32.mul(1.0, y), F32.mul(1.0, y))}]} : 447 Result<&2, &2, Backward.Error, List<&2, Maybe<&2, F32>>>} 448 449 # zero_grad empties every slot and keeps every node. 450 law zero_grad_empties: 451 for +entries: List<&2, Backward.Entry> 452 {tape_slots(Backward.zero_grad(Backward.Tape{entries})) == Backward.nones(List.length(&2, Backward.Entry, entries)) : List<&2, Maybe<&2, F32>>} 453 454 law zero_grad_keeps_nodes: 455 for +entries: List<&2, Backward.Entry> 456 {nodes_of(Backward.zero_grad(Backward.Tape{entries})) == Backward.nodes(entries) : List<&2, Backward.Node>}