def mapped_index

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

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/tensor/word_conversion.bend as Conversion

191def mapped_index(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +index: U32,
192  +mapping: Layout.Mapping<input, axes, plan, index>) ->
193  {Layout.source_index(input, axes, plan, index, mapping) ==
194    U32.from_nat(U32.to_nat(Maybe.default(&2, U32, Values.permute_source(input, axes, index), 0))) : U32}:
195  +found = Maybe.default(&2, U32, Values.permute_source(input, axes, index), 0)
196  Equal.trans(U32, Layout.source_index(input, axes, plan, index, mapping), found, U32.from_nat(U32.to_nat(found)),
197    Equal.cong(Maybe<&2, U32>, U32, result => Maybe.default(&2, U32, result, 0),
198      Some{Layout.source_index(input, axes, plan, index, mapping)}, Values.permute_source(input, axes, index),
199      reference(input, axes, plan, index, mapping)),
200    Equal.sym(U32, U32.from_nat(U32.to_nat(found)), found, Conversion.word_natural_exact(found)))