def mapping_reference

bendtorch-core@0.1.0.0 · src/tensor/permutation_mapping.bend:83

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/axis_algebra.bend as AxisAlgebra

83def mapping_reference(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Layout.Plan<input, axes>, +index: U32,
84  mapping: Layout.Mapping<input, axes, plan, index>) ->
85  {Some{Layout.source_index(input, axes, plan, index, mapping)} == Layout.reference_source(input, axes, plan, index) : Maybe<&2, U32>}:
86  Layout.Mapped{oi, oc, ic, source, requested, decoded, inverse, flattened, roundtrip} = mapping
87  Layout.CWord{+values, +valid} = ic
88  %AxisAlgebra.inverse_coordinates(input, axes,
89    Shape.words(Shape.digits(Shape.extents(Layout.output(input, axes, plan)), U32.to_nat(index)))) :
90      {Some{Layout.index_value(input, source)} == Layout.reference_source_coordinates(Shape.dimensions(input), _) : Maybe<&2, U32>}
91  %Equal.trans(List<&2, U32>, Layout.coordinate_values(Layout.output(input, axes, plan), oc),
92      Layout.reverse_words(Layout.unflatten_words(Layout.reverse_words(Shape.dimensions(Layout.output(input, axes, plan))), index)),
93      Shape.words(Shape.digits(Shape.extents(Layout.output(input, axes, plan)), U32.to_nat(index))), decoded,
94      codec_decoded(Shape.dimensions(Layout.output(input, axes, plan)), index)) :
95        {Some{Layout.index_value(input, source)} == Layout.reference_source_coordinates(Shape.dimensions(input),
96          Layout.inverse_coordinates(Shape.dimensions(input), axes, _, 0)) : Maybe<&2, U32>}
97  %inverse : {Some{Layout.index_value(input, source)} == Layout.reference_source_coordinates(Shape.dimensions(input), _) : Maybe<&2, U32>}
98  Equal.cong(U32, Maybe<&2, U32>, value => Some{value}, Layout.index_value(input, source),
99    U32.from_nat(Shape.offset(Shape.extents(input), Shape.naturals(values))),
100    Equal.trans(U32, Layout.index_value(input, source), Layout.flatten_words(Shape.dimensions(input), values, 0),
101      U32.from_nat(Shape.offset(Shape.extents(input), Shape.naturals(values))), flattened,
102      codec_flattened(Shape.dimensions(input), values, valid)))