def offset_shift

bendtorch-core@0.1.0.0 · src/autograd/reduce_parallel_proof.bend:355

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

355# The one word that rows of span 1 reach from `base` sits `base` past the
356# one they reach from 0.
357def offset_shift(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +target: Nat, +base: Nat) ->
358  {Reduce.single_offset(rows, extents, target, base) == Nat.add(base, Reduce.single_offset(rows, extents, target, 0n)) : Nat}:
359  match rows extents:
360    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others:
361      +inner = Nat.mod(target, left)
362      +step = Nat.mul(right, Nat.div(target, left))
363      +below = Reduce.single_offset(rest, others, inner, 0n)
364      Equal.trans(Nat, Reduce.single_offset(rest, others, inner, Nat.add(base, step)), Nat.add(Nat.add(base, step), below),
365        Nat.add(base, Reduce.single_offset(rest, others, inner, step)),
366        offset_shift(rest, others, inner, Nat.add(base, step)),
367        Equal.trans(Nat, Nat.add(Nat.add(base, step), below), Nat.add(base, Nat.add(step, below)),
368          Nat.add(base, Reduce.single_offset(rest, others, inner, step)),
369          Equal.sym(Nat, Nat.add(base, Nat.add(step, below)), Nat.add(Nat.add(base, step), below), Index.add_assoc(base, step, below)),
370          Equal.cong(Nat, Nat, x => Nat.add(base, x), Nat.add(step, below), Reduce.single_offset(rest, others, inner, step),
371            Equal.sym(Nat, Reduce.single_offset(rest, others, inner, step), Nat.add(step, below), offset_shift(rest, others, inner, step)))))
372    case Nil{} +extents:
373      add_zero(base)
374    case Shape.Row{+stride, +left, +right} <> +rest Nil{}:
375      add_zero(base)