def write_spec

bendtorch-core@0.1.0.0 · src/tensor/permute_build_proof.bend:623

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/tensor/contracts/build_facts.bend as Facts
import bendtorch-core@0.1.0.0/src/shape/index_proofs.bend as Index
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/storage/array.bend as ArrayCore

623# Writing cell k at the path whose position is k turns the spec for k into
624# the spec for 1 + k.
625def write_spec(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +k: Nat, +depth: Nat, +base: Nat, +path: Model.Path(depth),
626  +h: {Nat.add(Model.offset(depth, path), base) == k : Nat}) ->
627  {ArrayCore.write_path(F32, depth, spec(input, axes, fill, ops, k, depth, base), path, word(input, axes, fill, ops, k)) == spec(input, axes, fill, ops, 1n+k, depth, base) : Array<F32>}:
628  match depth:
629    case 0n:
630      %h : {ALeaf{word(input, axes, fill, ops, _)} == ALeaf{leaf(input, axes, fill, ops, base, Nat.is_lt(base, 1n+_))} : Array<F32>}
631      %Equal.sym(Bool, Nat.is_lt(base, 1n+base), True{}, Facts.lt_self_succ(base)) :
632        {ALeaf{word(input, axes, fill, ops, base)} == ALeaf{leaf(input, axes, fill, ops, base, _)} : Array<F32>}
633      {==}
634    case 1n+inner:
635      Model.Step{+bit, +rest} = path
636      match bit:
637        case False{}:
638          +c = Model.capacity(inner)
639          +high = Nat.add(c, base)
640          +inside = Equal.trans(Bool, Nat.is_lt(k, high), Nat.is_lt(Nat.add(Model.offset(inner, rest), base), high), True{},
641            Equal.cong(Nat, Bool, x => Nat.is_lt(x, high), k, Nat.add(Model.offset(inner, rest), base),
642              Equal.sym(Nat, Nat.add(Model.offset(inner, rest), base), k, h)),
643            Facts.lt_add_right(Model.offset(inner, rest), c, base, Api.offset_lt(inner, rest)))
644          %Equal.sym(Array<F32>, ArrayCore.write_path(F32, inner, spec(input, axes, fill, ops, k, inner, base), rest, word(input, axes, fill, ops, k)),
645            spec(input, axes, fill, ops, 1n+k, inner, base), write_spec(input, axes, fill, ops, k, inner, base, rest, h)) :
646            {ANode{_, spec(input, axes, fill, ops, k, inner, high)} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array<F32>}
647          %Equal.sym(Array<F32>, spec(input, axes, fill, ops, k, inner, high), spec(input, axes, fill, ops, 1n+k, inner, high), above(input, axes, fill, ops, k, inner, high, inside)) :
648            {ANode{spec(input, axes, fill, ops, 1n+k, inner, base), _} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array<F32>}
649          {==}
650        case True{}:
651          +c = Model.capacity(inner)
652          +high = Nat.add(c, base)
653          +off = Model.offset(inner, rest)
654          +lower = Equal.trans(Bool, Nat.is_le(high, k), Nat.is_le(high, Nat.add(c, Nat.add(off, base))), True{},
655            Equal.cong(Nat, Bool, x => Nat.is_le(high, x), k, Nat.add(c, Nat.add(off, base)),
656              Equal.trans(Nat, k, Nat.add(Nat.add(c, off), base), Nat.add(c, Nat.add(off, base)),
657                Equal.sym(Nat, Nat.add(Nat.add(c, off), base), k, h),
658                Equal.sym(Nat, Nat.add(c, Nat.add(off, base)), Nat.add(Nat.add(c, off), base), Index.add_assoc(c, off, base)))),
659            Facts.le_add_cong(c, base, Nat.add(off, base), Api.le_add_left(off, base)))
660          +shifted = Equal.trans(Nat, Nat.add(off, high), Nat.add(Nat.add(c, off), base), k, Facts.position_step(c, off, base), h)
661          %Equal.sym(Array<F32>, spec(input, axes, fill, ops, k, inner, base), spec(input, axes, fill, ops, 1n+k, inner, base), below(input, axes, fill, ops, k, inner, base, lower)) :
662            {ANode{_, ArrayCore.write_path(F32, inner, spec(input, axes, fill, ops, k, inner, high), rest, word(input, axes, fill, ops, k))} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array<F32>}
663          %Equal.sym(Array<F32>, ArrayCore.write_path(F32, inner, spec(input, axes, fill, ops, k, inner, high), rest, word(input, axes, fill, ops, k)),
664            spec(input, axes, fill, ops, 1n+k, inner, high), write_spec(input, axes, fill, ops, k, inner, high, rest, shifted)) :
665            {ANode{spec(input, axes, fill, ops, 1n+k, inner, base), _} == spec(input, axes, fill, ops, 1n+k, 1n+inner, base) : Array<F32>}
666          {==}