def expected_word

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
import bendtorch-core@0.1.0.0/src/tensor/observation.bend as Values
import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
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/refinement.bend as Refinement
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms

369# What the copy loop expects at output index `other`: the source word at
370# `permute_source_index`.
371def expected_word(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +other: Nat,
372  +ds: Nat, +ns: Nat, -a: Array<F32>, -b: Model.Balanced(F32, ds, a), +ok: {Storage.bounds(ds, ns) == True{} : Bool},
373  +fits: {ns == Shape.logical_numel(Shape.dimensions(input)) : Nat},
374  -source_fits: {Lifecycle.logical_length(Storage.Owned{ds, ns, a, b, ok}) == Shape.logical_numel(Shape.dimensions(input)) : Nat},
375  result: Result<&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, U32.from_nat(other)>>,
376  done: {Result.is_done(&2, &2, Layout.Error, Layout.Mapping<input, axes, plan, U32.from_nat(other)>, result) == True{} : Bool}) ->
377  Fact<{Transforms.expected_mapped(input, axes, plan, U32.from_nat(other), Core.TTensor{Storage.Owned{ds, ns, a, b, ok}, source_fits}, result) ==
378    Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}>:
379  match result:
380    case Fail{error}:
381      Fact{Empty.absurd({None{} == Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))} : Maybe<&2, F32>}, Arith.clash(done))}
382    case Done{+mapping}:
383      +found = Layout.source_index(input, axes, plan, U32.from_nat(other), mapping)
384      +inside : {Nat.is_lt(U32.to_nat(found), ns) == True{} : Bool} = %Equal.sym(Nat, ns, Shape.logical_numel(Shape.dimensions(input)), fits) : {Nat.is_lt(U32.to_nat(found), _) == True{} : Bool}
385        mapped_bound(input, axes, plan, U32.from_nat(other), mapping)
386      Fact{Equal.trans(Maybe<&2, F32>, Lifecycle.checked_observe_array(ds, ns, a, found), Model.observe(F32, ds, a, Refinement.route(ds, found)),
387        Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))},
388        owned_observe(ds, ns, a, found, ok, inside),
389        Equal.trans(Maybe<&2, F32>, Model.observe(F32, ds, a, Refinement.route(ds, found)), Some{Refinement.second(F32, Array.get(F32, a, found))},
390          Some{Refinement.second(F32, Array.get(F32, a, U32.from_nat(Values.permute_source_index(input, axes, other))))},
391          Equal.sym(Maybe<&2, F32>, Some{Refinement.second(F32, Array.get(F32, a, found))}, Model.observe(F32, ds, a, Refinement.route(ds, found)),
392            Refinement.law_get_refines(F32, ds, a, found, b)),
393          Equal.cong(U32, Maybe<&2, F32>, index => Some{Refinement.second(F32, Array.get(F32, a, index))}, found, U32.from_nat(Values.permute_source_index(input, axes, other)),
394            mapped_index(input, axes, plan, U32.from_nat(other), mapping))))}