def parallel_indexed

bendtorch-core@0.1.0.0 · src/autograd/reduce.bend:263

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise

263# An empty contribution leaves every target cell at +0.0.
264def parallel_indexed(input: Core.Dynamic, +target: Shape.Shape, +plan: Shape.BroadcastPlan, output: Core.Tensor<target>,
265  +depth: Nat, +length: Nat, -valid: {Storage.bounds(depth, length) == True{} : Bool},
266  -exact: {length == Shape.logical_numel(Shape.dimensions(target)) : Nat},
267  map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>) ->
268  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
269  match map:
270    case Shape.Unindexed{zero}:
271      (input, Done{Core.Packed{target, output}})
272    case Shape.Indexed{rows, same}:
273      parallel_rows(plan, target, depth, length, valid, exact, Elementwise.take(input), rows)