pymoose.select#

pymoose.select(x, axis, index, placement=None)[source]#

Select elements along some axis of a tensor according to some index tensor.

Parameters

  • x – A tensor.

  • axis – The dimension along which to index.

  • index – A 1-d boolean tensor such that len(index) == x.shape[axis].

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

Returns

Copy of the tensor x, with the values along axis filtered such that x[..., i, ...] is kept iff index[i] == 1.