pymoose.index_axis

pymoose.index_axis(x, axis, index, placement=None)

Index a tensor along a given dimension.

Parameters

• x – A tensor.

• axis – The dimension along which to index.

• placement – An optional Placement to pin this operation to.

Returns

A slice of x corresponding to the sub-tensor at index along axis.