masked_fill¶

nvtripy.masked_fill(input: Tensor, mask: Tensor, value: Number) → Tensor[source]¶

Returns a new tensor filled with value where mask is True and elements from the input tensor otherwise.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
mask (Tensor) – [dtype=T2] The mask tensor.
value (Number) – the value to fill with. This will be casted to match the data type of the input tensor.

Returns:

[dtype=T1] A new tensor of the same shape as the input tensor.

Return type:

Tensor

DATA TYPE CONSTRAINTS:

T1: bool, float16, bfloat16, float32, int4, int8, int32, int64
T2: bool

Example

mask = tp.Tensor([[True, False], [True, True]])
input = tp.zeros([2, 2])
output = tp.masked_fill(input, mask, -1.0)

Local Variables¶

>>> mask
tensor(
    [[True, False],
     [True, True]], 
    dtype=bool, loc=cpu:0, shape=(2, 2))

>>> input
tensor(
    [[0, 0],
     [0, 0]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))

>>> output
tensor(
    [[-1, 0],
     [-1, -1]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))