def fiber_eq

bendtorch-core@0.1.0.0 · src/autograd/reduce_replayed_proof.bend:468

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/tensor/contracts/matmul_parallel_proof.bend as Par
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
import bendtorch-core@0.1.0.0/src/autograd/reduce_parallel_proof.bend as ReduceParallel

468def fiber_eq(+rows: List<&2, Shape.Row>, +extents: List<&2, Nat>, +fuel: Nat, +target: Nat, +base: Nat, +digit: Nat, +accumulator: F32,
469  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +d: Nat, +f: F32,
470  +source: Par.Source<shape, fill, ops, h, d, f>,
471  +below: {ReduceParallel.all_below(Fiber.visited(rows, extents, fuel, target, base, digit), Shape.logical_numel(Shape.dimensions(shape))) == True{} : Bool}) ->
472  {Reduce.fiber(rows, extents, fuel, target, base, digit, accumulator, Api.arr(h, d, f)) ==
473    (rfold(shape, fill, ops, Fiber.visited(rows, extents, fuel, target, base, digit), accumulator), Api.arr(h, d, f)) : F32 & Array<F32>}:
474  match rows extents fuel:
475    case Nil{} _ _:
476      leaf_eq(base, accumulator, shape, fill, ops, h, d, f, source,
477        Arith.and_left(Nat.is_lt(base, Shape.logical_numel(Shape.dimensions(shape))), True{}, below),
478        Array.get(F32, Api.arr(h, d, f), U32.from_nat(base)), {==})
479    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 0n:
480      {==}
481    case Shape.Row{+stride, +left, +right} <> +rest +extent <> +others 1n+remaining:
482      +rows : List<&2, Shape.Row> = Shape.Row{stride, left, right} <> rest
483      +extents : List<&2, Nat> = extent <> others
484      +inner = Nat.mod(target, left)
485      +offset = Nat.add(base, Nat.mul(right, digit))
486      +low = Fiber.visited(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner))
487      +high = Fiber.visited(rows, extents, remaining, target, base, 1n+digit)
488      +bound = Shape.logical_numel(Shape.dimensions(shape))
489      +value = rfold(shape, fill, ops, low, accumulator)
490      Equal.trans(F32 & Array<F32>,
491        Reduce.walk_next(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k),
492          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f))),
493        Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, value, Api.arr(h, d, f)),
494        (rfold(shape, fill, ops, List.append(&2, Nat, low, high), accumulator), Api.arr(h, d, f)),
495        next_eq(v => k => Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, v, k),
496          Reduce.fiber(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator, Api.arr(h, d, f)),
497          value, Api.arr(h, d, f),
498          fiber_eq(rest, others, Reduce.first_span(rest, others), inner, offset, Reduce.first_digit(rest, inner), accumulator,
499            shape, fill, ops, h, d, f, source, below_front(low, high, bound, below))),
500        %Equal.sym(F32, rfold(shape, fill, ops, List.append(&2, Nat, low, high), accumulator), rfold(shape, fill, ops, high, value),
501          rfold_append(shape, fill, ops, low, high, accumulator)) :
502          {Reduce.fiber(rows, extents, remaining, target, base, 1n+digit, value, Api.arr(h, d, f)) == (_, Api.arr(h, d, f)) : F32 & Array<F32>}
503        fiber_eq(rows, extents, remaining, target, base, 1n+digit, value, shape, fill, ops, h, d, f, source, below_back(low, high, bound, below)))
504    case Shape.Row{+stride, +left, +right} <> +rest Nil{} _:
505      {==}