def map_roundtrip

bendtorch-core@0.1.0.0 · src/tensor/layout.bend:336

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape

336def map_roundtrip(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
337  oi: FlatIndex<output(input, axes, plan)>, +oc: CheckedCoordinate<output(input, axes, plan)>,
338  +ic: CheckedCoordinate<input>,
339  requested: {index_value(output(input, axes, plan), oi) == i : U32},
340  decoded: {coordinate_values(output(input, axes, plan), oc) ==
341    reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
342  inverse: {Some{coordinate_values(input, ic)} ==
343    inverse_coordinates(Shape.dimensions(input), axes, coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
344  decision: Shape.BoolCase<maybe_same(reorder(coordinate_values(input, ic), axes),
345    Some{coordinate_values(output(input, axes, plan), oc)})>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
346  match decision:
347    case Shape.BFalse{equal}:
348      Fail{MappingError{}}
349    case Shape.BTrue{equal}:
350      +j = flatten_words(Shape.dimensions(input), coordinate_values(input, ic), 0)
351      map_flat(input, axes, plan, i, oi, oc, ic, requested, decoded, inverse, equal, j, {==},
352        Shape.bool_case(U32.is_lt(j, Shape.count(input))))