def fiber_walk

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

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

846# The fiber-walk lemma: for the index rows of an aligned broadcast and a
847# cell below the left count, the offsets the walk reads, in order, are the
848# masked positions.
849def fiber_walk(+left: List<&2, Nat>, +right: List<&2, Nat>, +output: List<&2, Nat>, +aligned: Shape.Aligned(left, right, output),
850  +target: Nat, +below: {Nat.is_lt(target, Shape.numel(left)) == True{} : Bool}) -> RowsWalk(left, right, output, target):
851  match left right output:
852    case Nil{} Nil{} Nil{}:
853      %Arith.lt_one(target, below) : {kept(Nat.is_eq(0n, _), 0n, Nil{}) == [0n] : List<&2, Nat>}
854      {==}
855    case +lh <> +lt +rh <> +rt +oh <> +ot:
856      Shape.AStep{+axis, +tail} = aligned
857      fiber_cons(lh, lt, rh, rt, oh, ot, axis, tail, target, below,
858        inner => inside => fiber_walk(lt, rt, ot, tail, inner, inside))
859    case Nil{} Nil{} +oh <> +ot:
860      Empty.absurd(RowsWalk(Nil{}, Nil{}, oh <> ot, target), aligned)
861    case Nil{} +rh <> +rt Nil{}:
862      Empty.absurd(RowsWalk(Nil{}, rh <> rt, Nil{}, target), aligned)
863    case +lh <> +lt Nil{} Nil{}:
864      Empty.absurd(RowsWalk(lh <> lt, Nil{}, Nil{}, target), aligned)
865    case Nil{} +rh <> +rt +oh <> +ot:
866      Empty.absurd(RowsWalk(Nil{}, rh <> rt, oh <> ot, target), aligned)
867    case +lh <> +lt Nil{} +oh <> +ot:
868      Empty.absurd(RowsWalk(lh <> lt, Nil{}, oh <> ot, target), aligned)
869    case +lh <> +lt +rh <> +rt Nil{}:
870      Empty.absurd(RowsWalk(lh <> lt, rh <> rt, Nil{}, target), aligned)