def map_flat

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

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

318def map_flat(+input: Shape.Shape, +axes: List<&2, U32>, +plan: Plan<input, axes>, +i: U32,
319  oi: FlatIndex<output(input, axes, plan)>, oc: CheckedCoordinate<output(input, axes, plan)>,
320  +ic: CheckedCoordinate<input>,
321  requested: {index_value(output(input, axes, plan), oi) == i : U32},
322  decoded: {coordinate_values(output(input, axes, plan), oc) ==
323    reverse_words(unflatten_words(reverse_words(Shape.dimensions(output(input, axes, plan))), i)) : List<&2, U32>},
324  inverse: {Some{coordinate_values(input, ic)} ==
325    inverse_coordinates(Shape.dimensions(input), axes, coordinate_values(output(input, axes, plan), oc), 0) : Maybe<&2, List<&2, U32>>},
326  roundtrip: {True{} == maybe_same(reorder(coordinate_values(input, ic), axes),
327    Some{coordinate_values(output(input, axes, plan), oc)}) : Bool},
328  +j: U32, flat: {j == flatten_words(Shape.dimensions(input), coordinate_values(input, ic), 0) : U32},
329  decision: Shape.BoolCase<U32.is_lt(j, Shape.count(input))>) -> Result<&2, &2, Error, Mapping<input, axes, plan, i>>:
330  match decision:
331    case Shape.BFalse{equal}:
332      Fail{MappingError{}}
333    case Shape.BTrue{equal}:
334      Done{Mapped{oi, oc, ic, FIndex{j, equal}, requested, decoded, inverse, flat, roundtrip}}