def chain

bendtorch-core@0.1.0.0 · src/tensor/permute_values.bend:90

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/word_conversion.bend as Conversion
import bendtorch-core@0.1.0.0/src/tensor/from_values_contract.bend as FromValues
import bendtorch-core@0.1.0.0/src/tensor/index_map.bend as IndexMap
import bendtorch-core@0.1.0.0/src/tensor/permute_build_proof.bend as PermuteProof
import bendtorch-core@0.1.0.0/src/tensor/contracts/build_facts.bend as Facts
import bendtorch-core@0.1.0.0/src/tensor/contracts/binary_storage_proof.bend as Owned
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/storage/array.bend as ArrayCore
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
import bendtorch-core@0.1.0.0/src/storage/certificates.bend as Certificates

90# The copy loop's write history, replayed on the owned storage that
91# `created` makes, is the owned storage that holds the output spec.
92def chain(+input: Shape.Shape, +axes: List<&2, U32>, +fill: F32, +ops: List<&2, StorageLaws.Op>, +d: Nat, +n: Nat, -f: {Storage.bounds(d, n) == True{} : Bool}, +fits: {Storage.bounds(d, n) == True{} : Bool},
93  +limit: U32, +n_limit: {Nat.is_le(n, U32.to_nat(limit)) == True{} : Bool},
94  +k: Nat, +room: {Nat.is_le(k, n) == True{} : Bool}) ->
95  {StorageLaws.replay(IndexMap.writes(~Request, ~source_of, k, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}) == Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, k, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, k, d, 0n), f} : Storage.Storage}:
96  match k:
97    case 0n:
98      Owned.owned_eq(d, n, Array.new(F32, d, 0.0), PermuteProof.spec(input, axes, fill, ops, 0n, d, 0n), PermuteProof.new_spec(input, axes, fill, ops, d, 0n), ArrayCore.new_balanced(F32, d, 0.0), f)(
99        PermuteProof.spec_balanced(input, axes, fill, ops, 0n, d, 0n))
100    case 1n+j:
101      +index = U32.from_nat(j)
102      +value = PermuteProof.word(input, axes, fill, ops, j)
103      +below_n = Conversion.le_successor_lt(j, n, room)
104      +exact = Conversion.natural_word_exact(j, limit, Facts.le_trans(j, n, U32.to_nat(limit), Conversion.lt_le(j, n, below_n), n_limit))
105      +inside = Equal.trans(Bool, Nat.is_lt(U32.to_nat(index), n), Nat.is_lt(j, n), True{},
106        Equal.cong(Nat, Bool, x => Nat.is_lt(x, n), U32.to_nat(index), j, exact), below_n)
107      +offset = Equal.trans(Nat, Nat.add(Model.offset(d, Refinement.route(d, index)), 0n), U32.to_nat(index), j,
108        Facts.output_position(d, n, index, fits, inside), exact)
109      +written = Equal.trans(Array<F32>, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), ArrayCore.write_path(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), Refinement.route(d, index), value),
110        PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), Refinement.law_set_refines(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n)),
111        PermuteProof.write_spec(input, axes, fill, ops, j, d, 0n, Refinement.route(d, index), offset))
112      Equal.trans(Storage.Storage, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, 1n+j, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
113        StorageLaws.step(StorageLaws.replay(IndexMap.writes(~Request, ~source_of, j, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}), StorageLaws.Write{index, value}),
114        Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n), f},
115        FromValues.replay_snoc(IndexMap.writes(~Request, ~source_of, j, Request{input, axes}, input, fill, ops), StorageLaws.Write{index, value}, Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}),
116        %Equal.sym(Storage.Storage, StorageLaws.replay(IndexMap.writes(~Request, ~source_of, j, Request{input, axes}, input, fill, ops), Storage.Owned{d, n, Array.new(F32, d, 0.0), ArrayCore.new_balanced(F32, d, 0.0), f}), Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n), f},
117          chain(input, axes, fill, ops, d, n, f, fits, limit, n_limit, j, Conversion.le_predecessor(j, n, room))) :
118          {StorageLaws.step(_, StorageLaws.Write{index, value}) == Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n), f} : Storage.Storage}
119        %Equal.sym(Storage.Storage & Result<&2, &2, Storage.Error, Unit>,
120          Storage.set_checked(d, n, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Nat.is_lt(U32.to_nat(index), n), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n), f),
121          (Storage.Owned{d, n, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value),
122            Certificates.set_result(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n)), f}, Done{Unit{}}),
123          Api.set_checked_hit(d, n, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Nat.is_lt(U32.to_nat(index), n), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n), f,
124            Equal.sym(Bool, Nat.is_lt(U32.to_nat(index), n), True{}, inside), fits, inside)) :
125          {StorageLaws.owner(Result<&2, &2, Storage.Error, Unit>, _) == Storage.Owned{d, n, PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n), f} : Storage.Storage}
126        Owned.owned_eq(d, n, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), PermuteProof.spec(input, axes, fill, ops, 1n+j, d, 0n), written,
127          Certificates.set_result(F32, d, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value, Array.set(F32, PermuteProof.spec(input, axes, fill, ops, j, d, 0n), index, value), {==}, PermuteProof.spec_balanced(input, axes, fill, ops, j, d, 0n)), f)(
128          PermuteProof.spec_balanced(input, axes, fill, ops, 1n+j, d, 0n)))