def flat_input_indices

bendtorch-core@0.1.0.0 · src/tensor/elementwise.bend:120

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/flat.bend as Flat
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core

120def flat_input_indices(+plan: Shape.BroadcastPlan,
121  result: Result<&2, &2, Shape.Error, Shape.FlatIndex<Shape.output_shape(plan)>>) -> Result<&2, &1, Core.Error, U32 & U32>:
122  match result:
123    case Fail{error}:
124      Fail{Core.ShapeError{error}}
125    case Done{+index}:
126      +map = Shape.index_map(plan)
127      Done{(Shape.index_value(Shape.left_shape(plan), Flat.project_left_index(plan, map, index)),
128        Shape.index_value(Shape.right_shape(plan), Flat.project_right_index(plan, map, index)))}