def reduced_blocks

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

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

662# A reduced row (left stride 0) walks every digit: the blocks digit ..
663# digit + f - 1 of the masked positions are what the walk visits.
664def reduced_blocks(+f: Nat, +stride: Nat, +right: Nat, +rest: List<&2, Shape.Row>, +extent: Nat, +others: List<&2, Nat>,
665  +target: Nat, +base: Nat, +digit: Nat, +walk: RestWalk(rest, others, target, stride)) ->
666  {masked(Nat.mul(f, stride), Shape.Row{stride, 0n, right} <> rest, target, base, Nat.mul(digit, stride)) ==
667    visited(Shape.Row{stride, 0n, right} <> rest, extent <> others, f, target, base, digit) : List<&2, Nat>}:
668  match f:
669    case 0n:
670      {==}
671    case 1n+g:
672      +rows : List<&2, Shape.Row> = Shape.Row{stride, 0n, right} <> rest
673      +c = Nat.mul(digit, stride)
674      +w = visited(rest, others, Reduce.first_span(rest, others), target, Nat.add(base, Nat.mul(right, digit)), Reduce.first_digit(rest, target))
675      +v = visited(rows, extent <> others, g, target, base, 1n+digit)
676      %Equal.sym(List<&2, Nat>, masked(Nat.add(stride, Nat.mul(g, stride)), rows, target, base, c),
677        List.append(&2, Nat, masked(stride, rows, target, base, c), masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))),
678        masked_split(stride, Nat.mul(g, stride), rows, target, base, c)) :
679        {_ == List.append(&2, Nat, w, v) : List<&2, Nat>}
680      %Equal.sym(List<&2, Nat>, masked(stride, rows, target, base, c), masked(stride, rest, target, Nat.add(base, Nat.mul(right, digit)), 0n),
681        block(stride, 0n, right, rest, target, target, base, digit, Free{{==}, {==}})) :
682        {List.append(&2, Nat, _, masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))) == List.append(&2, Nat, w, v) : List<&2, Nat>}
683      %Equal.sym(List<&2, Nat>, masked(stride, rest, target, Nat.add(base, Nat.mul(right, digit)), 0n), w,
684        rest_at(rest, others, target, stride, walk, Nat.add(base, Nat.mul(right, digit)))) :
685        {List.append(&2, Nat, _, masked(Nat.mul(g, stride), rows, target, base, Nat.add(c, stride))) == List.append(&2, Nat, w, v) : List<&2, Nat>}
686      %next_cursor(digit, stride) :
687        {List.append(&2, Nat, w, masked(Nat.mul(g, stride), rows, target, base, _)) == List.append(&2, Nat, w, v) : List<&2, Nat>}
688      %Equal.sym(List<&2, Nat>, masked(Nat.mul(g, stride), rows, target, base, Nat.mul(1n+digit, stride)), v,
689        reduced_blocks(g, stride, right, rest, extent, others, target, base, 1n+digit, walk)) :
690        {List.append(&2, Nat, w, _) == List.append(&2, Nat, w, v) : List<&2, Nat>}
691      {==}