def tree_eq

bendtorch-core@0.1.0.0 · src/tensor/contracts/binary_build_proof.bend:267

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/main.bend as Storage
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel

267def tree_eq(+operation: Elementwise.Operation, +plan: Shape.BroadcastPlan, +count: Nat, +ha: List<&2, Api.Effect>, +dl: Nat, +fa: F32, +nl: Nat, +hb: List<&2, Api.Effect>, +dr: Nat, +fb: F32, +nr: Nat, +lfits: {Storage.bounds(dl, nl) == True{} : Bool}, +rfits: {Storage.bounds(dr, nr) == True{} : Bool}, +lvalid: {written(nl, ha) == True{} : Bool}, +rvalid: {written(nr, hb) == True{} : Bool}, +count_exact: {count == Shape.logical_numel(Shape.dimensions(Shape.output_shape(plan))) : Nat}, +left_count: {nl == Shape.logical_numel(Shape.dimensions(Shape.left_shape(plan))) : Nat}, +right_count: {nr == Shape.logical_numel(Shape.dimensions(Shape.right_shape(plan))) : Nat}, +depth: Nat, +base: Nat) ->
268  TreeEq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, depth, base, Parallel.tree(~Elementwise.Walk, ~Array<F32>, ~Array<F32>, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), depth, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)))):
269  match depth:
270    case 0n:
271      cell_at_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, Nat.is_lt(base, count), {==})
272    case 1n+inner:
273      match inner:
274        case 0n:
275          pair_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, Elementwise.cell_at(Elementwise.walk(Elementwise.Request{operation, plan}), base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb)), Nat.is_lt(base, count)),
276            cell_at_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, base, Nat.is_lt(base, count), {==}))
277        case 1n+below:
278          tree_node_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, inner, base, Parallel.tree(~Elementwise.Walk, ~Array<F32>, ~Array<F32>, ~Elementwise.cell_at, Elementwise.walk(Elementwise.Request{operation, plan}), inner, count, base, (Api.arr(ha, dl, fa), Api.arr(hb, dr, fb))), tree_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, inner, base),
279            tree_eq(operation, plan, count, ha, dl, fa, nl, hb, dr, fb, nr, lfits, rfits, lvalid, rvalid, count_exact, left_count, right_count, inner, Nat.add(Model.capacity(inner), base)))