torch_npu.npu_masked_fill_range

API接口

torch_npu.npu_masked_fill_range(self, start, end, value, axis=-1) -> Tensor

功能描述

同轴上被range.boxes屏蔽（masked）的填充张量。自定义屏蔽填充范围算子。

参数说明

• self (Tensor) - shape为1D (D,)、2D (N,D)或3D (N,D)的float32/float16/int32/int8 ND张量。
• start (Tensor) - 屏蔽填充开始位置。shape为(num,N)的int32 3D张量。
• end (Tensor) - 屏蔽填充结束位置。shape为(num,N)的int32 3D张量。
• value (Tensor) - 屏蔽填充值。shape为(num,)的float32/float16/int32/int8 2D张量。
• axis (Int，默认值为-1) - 带有int32屏蔽填充的轴。

示例

>>> a=torch.rand(4,4).npu()
>>> a
tensor([[0.9419, 0.4919, 0.2874, 0.6560],
        [0.6691, 0.6668, 0.0330, 0.1006],
        [0.3888, 0.7011, 0.7141, 0.7878],
        [0.0366, 0.9738, 0.4689, 0.0979]], device='npu:0')
>>> start = torch.tensor([[0,1,2]], dtype=torch.int32).npu()
>>> end = torch.tensor([[1,2,3]], dtype=torch.int32).npu()
>>> value = torch.tensor([1], dtype=torch.float).npu()
>>> out = torch_npu.npu_masked_fill_range(a, start, end, value, 1)
>>> out
tensor([[1.0000, 0.4919, 0.2874, 0.6560],
        [0.6691, 1.0000, 0.0330, 0.1006],
        [0.3888, 0.7011, 1.0000, 0.7878],
        [0.0366, 0.9738, 0.4689, 0.0979]], device='npu:0')