oneflow.masked_select

oneflow.masked_select(input, mask)

Returns a new 1-D tensor which indexes the input tensor according to the boolean mask mask which is a BoolTensor(In oneFlow BoolTensor is replaced by Int8Tensor).

The shapes of the mask tensor and the input tensor don’t need to match, but they must be broadcastable.

Parameters
  • input (Tensor) – the input tensor.

  • mask (Tensor) – the tensor containing the binary mask to index with

For example:

>>> import oneflow as flow
>>> import numpy as np

>>> input = flow.tensor(np.array([[-0.4620, 0.3139], [0.3898, -0.7197], [0.0478, -0.1657]]), dtype=flow.float32)
>>> mask = input.gt(0.05)
>>> out = flow.masked_select(input, mask)
>>> out
tensor([0.3139, 0.3898], dtype=oneflow.float32)