def scatter_gather

bendtorch-core@0.1.0.0 · src/tensor/axis_algebra.bend:3085

import bendtorch-core@0.1.0.0/src/tensor/axis_algebra.bend as AxisAlgebra
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_spec.bend as Axis

3085def scatter_gather(+shape: Shape.Shape, +axes: List<&2, U32>, +coords: List<&2, U32>,
3086  accepted: {True{} == Layout.bijective(shape, axes) : Bool}, length: {Shape.rank_dimensions(coords) == Shape.rank_dimensions(Shape.dimensions(shape)) : Nat}) ->
3087  {Layout.inverse_coordinates(Shape.dimensions(shape), axes, Axis.order(U32, coords, Shape.naturals(axes)), 0) == Some{coords} : Maybe<&2, List<&2, U32>>}:
3088  certificate = facts_lists(U32, Shape.rank_dimensions(Shape.dimensions(shape)), Shape.naturals(axes), coords,
3089    Equal.trans(Nat, List.length(&2, U32, coords), Shape.rank_dimensions(coords), Shape.rank_dimensions(Shape.dimensions(shape)), word_length(coords), length), shape_facts(shape, axes, accepted))
3090  Equal.trans(Maybe<&2, List<&2, U32>>, Layout.inverse_coordinates(Shape.dimensions(shape), axes, Axis.order(U32, coords, Shape.naturals(axes)), 0),
3091    Axis.scatter(U32, Shape.rank_dimensions(Shape.dimensions(shape)), Shape.naturals(axes), Axis.order(U32, coords, Shape.naturals(axes)), 0n), Some{coords},
3092    inverse_coordinates(shape, axes, Axis.order(U32, coords, Shape.naturals(axes))), scatter_restored(U32, Shape.rank_dimensions(Shape.dimensions(shape)), Shape.naturals(axes), coords, certificate))

Uses (20)

Used by (1)