def masked_split

bendtorch-core@0.1.0.0 · src/autograd/reduce_fiber_proof.bend:277

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/broadcast_proofs.bend as Broadcast
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith

277def masked_split(+a: Nat, +b: Nat, +rows: List<&2, Shape.Row>, +target: Nat, +base: Nat, +cursor: Nat) ->
278  {masked(Nat.add(a, b), rows, target, base, cursor) ==
279    List.append(&2, Nat, masked(a, rows, target, base, cursor), masked(b, rows, target, base, Nat.add(cursor, a))) : List<&2, Nat>}:
280  match a:
281    case 0n:
282      %Equal.sym(Nat, Nat.add(cursor, 0n), cursor, Broadcast.add_zero(cursor)) :
283        {masked(b, rows, target, base, cursor) == masked(b, rows, target, base, _) : List<&2, Nat>}
284      {==}
285    case 1n+p:
286      +keep = Nat.is_eq(Shape.walk_left(rows, cursor), target)
287      +position = Nat.add(base, Shape.walk_right(rows, cursor))
288      +front = masked(p, rows, target, base, 1n+cursor)
289      %Equal.sym(Nat, Nat.add(cursor, 1n+p), 1n+Nat.add(cursor, p), Arith.add_succ(cursor, p)) :
290        {masked(1n+Nat.add(p, b), rows, target, base, cursor) ==
291          List.append(&2, Nat, kept(keep, position, front), masked(b, rows, target, base, _)) : List<&2, Nat>}
292      Equal.trans(List<&2, Nat>, kept(keep, position, masked(Nat.add(p, b), rows, target, base, 1n+cursor)),
293        kept(keep, position, List.append(&2, Nat, front, masked(b, rows, target, base, Nat.add(1n+cursor, p)))),
294        List.append(&2, Nat, kept(keep, position, front), masked(b, rows, target, base, Nat.add(1n+cursor, p))),
295        Equal.cong(List<&2, Nat>, List<&2, Nat>, xs => kept(keep, position, xs),
296          masked(Nat.add(p, b), rows, target, base, 1n+cursor),
297          List.append(&2, Nat, front, masked(b, rows, target, base, Nat.add(1n+cursor, p))),
298          masked_split(p, b, rows, target, base, 1n+cursor)),
299        kept_append(keep, position, front, masked(b, rows, target, base, Nat.add(1n+cursor, p))))