def fiber_cons

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

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

822def fiber_cons(+lh: Nat, +lt: List<&2, Nat>, +rh: Nat, +rt: List<&2, Nat>, +oh: Nat, +ot: List<&2, Nat>, +axis: Shape.Axis<lh, rh, oh>,
823  +tail: Shape.Aligned(lt, rt, ot), +target: Nat, +below: {Nat.is_lt(target, Nat.mul(lh, Shape.numel(lt))) == True{} : Bool},
824  rec: @+inner: Nat -> @+inside: {Nat.is_lt(inner, Shape.numel(lt)) == True{} : Bool} -> RowsWalk(lt, rt, ot, inner)) ->
825  RowsWalk(lh <> lt, rh <> rt, oh <> ot, target):
826  match lh:
827    case 0n:
828      Empty.absurd(RowsWalk(0n <> lt, rh <> rt, oh <> ot, target), Arith.lt_zero(target, below))
829    case 1n+k:
830      match k:
831        case 0n:
832          reduced_blocks(oh, Shape.numel(ot), Shape.broadcast_index(rh, Shape.numel(rt)), Shape.index_rows(lt, rt, ot), oh, ot, target, 0n, 0n,
833            rec(target,
834              %Broadcast.add_zero(Shape.numel(lt)) : {Nat.is_lt(target, _) == True{} : Bool}
835              below))
836        case 1n+j:
837          +size = Shape.numel(lt)
838          kept_row(size, Shape.numel(ot), Shape.broadcast_index(rh, Shape.numel(rt)), Shape.index_rows(lt, rt, ot), oh, ot, target, 0n,
839            %Equal.sym(Nat, oh, 2n+j, axis_kept(j, rh, oh, axis)) : {Nat.is_lt(target, Nat.mul(_, size)) == True{} : Bool}
840            below,
841            rec(Nat.mod(target, size),
842              {\{Decomposed: m_below => sum => m_below} : Decomposed<target, size> -> {Nat.is_lt(Nat.mod(target, size), size) == True{} : Bool}}(
843                decompose(target, size, positive_of(target, 2n+j, size, below)))),
844            BoundBy{lt, rt, ot, tail, {==}, {==}, {==}})