Contents

gatherAlongAxis(indices:axis:)

Adds a gather-along-axis operation to the current graph.

Declaration

func gatherAlongAxis(indices: some BNNSGraph.Builder.OperationParameter<Int32>, axis: Int) -> BNNSGraph.Builder.Tensor<T>

Parameters

  • indices:

    The indices that the operation gathers.

  • axis:

    The axis over which the operation gathers.

Discussion

Register a gather along axis op, which shuffles the input along one axis. In particular, output y such that `y[L…, i, T…] = x[L…, indices[L…, i, T…], T…].