昇腾社区首页
中文
注册

torch_npu.npu_broadcast

API接口

torch_npu.npu_broadcast(self, size) -> Tensor

功能描述

返回self张量的新视图,其单维度扩展,结果连续。

张量也可以扩展更多维度,新的维度添加在最前面。

参数说明

  • self (Tensor) - 输入张量。
  • size (ListInt) - 对应扩展尺寸。

示例

>>> x = torch.tensor([[1], [2], [3]]).npu()
>>> x.shape
torch.Size([3, 1])
>>> x.npu_broadcast(3, 4)
tensor([[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3]], device='npu:0')