def build_eq

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
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/tensor/parallel.bend as Parallel
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

877# Whatever the split, the parallel build over the copy returns the spec.
878def build_eq(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
879  +rows_exact: {rows == Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))) : List<&2, Shape.Row>},
880  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +d: Nat, +f: F32, +count: Nat,
881  +source: Par.Source<shape, fill, ops, h, d, f>, +facts: Facts<plan, shape, count>, +depth: Nat, +splits: Nat, +base: Nat) ->
882  {Reduce.build(depth, splits, rows, Shape.extents(Shape.output_shape(plan)), count, base, Api.arr(h, d, f)) ==
883    out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, depth, base) : Array<F32>}:
884  match depth splits:
885    case 1n+lower 1n+rest:
886      fork_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, lower, rest, base,
887        build_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, lower, rest, base),
888        build_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, lower, rest, Nat.add(Model.capacity(lower), base)),
889        Array.clone(F32, Api.arr(h, d, f)), {==})
890    case 0n _:
891      tree_first(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 0n, base,
892        Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), 0n, count, base, (Api.arr(h, d, f), Unit{})),
893        tree_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, 0n, base))
894    case 1n+lower 0n:
895      tree_first(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 1n+lower, base,
896        Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), 1n+lower, count, base, (Api.arr(h, d, f), Unit{})),
897        tree_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, 1n+lower, base))