def fiber_run

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

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/tensor/elementwise_parallel_proof.bend as Parallel
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
import bendtorch-core@0.1.0.0/src/autograd/reduce_fiber_proof.bend as Fiber

318def fiber_run(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +fuel: Nat, +target: Nat, +base: Nat, +digit: Nat, +d: Nat,
319  -A: Array<F32>, +accumulator: F32, array: Array<F32>, -same: {array == A : Array<F32>}, -valid: Model.Balanced(F32, d, A)) ->
320  Parallel.Fact<FiberRun<d, A, Fiber.visited(rows, extents, fuel, target, base, digit), accumulator,
321    Reduce.fiber(rows, extents, fuel, target, base, digit, accumulator, array)>>:
322  match rows extents fuel:
323    case Nil{} _ _:
324      leaf_run(d, A, base, accumulator, array, Array.get(F32, array, U32.from_nat(base)), {==}, same, valid)
325    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
326      Parallel.Fact{FiberRun{Nil{}, Unit{}, {==}, same}}
327    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
328      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
329      +extents : List<&2, Nat> = extent <> others
330      +inner = Nat.mod(target, left)
331      next_run(d, A,
332        Fiber.visited(rest, others, Reduce.first_span(rest, others), inner, Nat.add(base, Nat.mul(right, digit)), Reduce.first_digit(rest, inner)),
333        Fiber.visited(rows, extents, remaining, target, base, 1n+digit), accumulator,
334        value => kept => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, value, kept),
335        Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, Nat.add(base, Nat.mul(right, digit)), Reduce.first_digit(rest, inner),
336          accumulator, array),
337        fiber_run(rest, others, Reduce.first_span(rest, others), inner, Nat.add(base, Nat.mul(right, digit)), Reduce.first_digit(rest, inner),
338          d, A, accumulator, array, same, valid),
339        value => kept => same => fiber_run(rows, extents, remaining, target, base, 1n+digit, d, A, value, kept, same, valid))
340    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
341      Parallel.Fact{FiberRun{Nil{}, Unit{}, {==}, same}}