支持DTensor的API清单
昇腾支持DTensor的API列表请参见表1。
Note
- DTensor（Distributed Tensor）是PyTorch提供的分布式张量抽象，是
torch.Tensor的子类。DTensor遵循SPMD（Single Program, Multiple Data）编程模型，通过DeviceMesh和Placement（包括Shard、Replicate、Partial）描述张量在多设备上的分片布局，透明地处理分片存储、算子计算和跨设备集合通信，使用户能够像编写单设备程序一样编写分布式程序。详细信息请参见PyTorch DTensor官方文档。 - 本文档展示的列表为已验证过支持DTensor的API。部分未验证的API可能也具备支持DTensor的能力，在昇腾NPU环境中请谨慎使用。我们将持续验证更多API并定期更新本文档。
| API名称 |
|---|
| Tensor.add |
| Tensor.add_ |
| Tensor.addcdiv_ |
| Tensor.addcmul_ |
| Tensor.addmm |
| Tensor.all |
| Tensor.any |
| Tensor.bitwise_and |
| Tensor.bitwise_not |
| Tensor.bitwise_or |
| Tensor.bmm |
| Tensor.clamp |
| Tensor.copy_ |
| Tensor.cos |
| Tensor.cumsum |
| Tensor.div |
| Tensor.div_ |
| Tensor.eq |
| Tensor.exponential_ |
| Tensor.fill_ |
| Tensor.flatten |
| Tensor.gather |
| Tensor.ge |
| Tensor.gt |
| Tensor.histc |
| Tensor.index_select |
| Tensor.isnan |
| Tensor.le |
| Tensor.log2 |
| Tensor.logical_and_ |
| Tensor.logical_not |
| Tensor.lt |
| Tensor.masked_fill |
| Tensor.matmul |
| Tensor.max |
| Tensor.maximum |
| Tensor.mean |
| Tensor.min |
| Tensor.minimum |
| Tensor.mm |
| Tensor.mul |
| Tensor.mul_ |
| Tensor.neg |
| Tensor.permute |
| Tensor.pow |
| Tensor.reciprocal |
| Tensor.remainder |
| Tensor.repeat |
| Tensor.reshape |
| Tensor.rsqrt |
| Tensor.scatter_ |
| Tensor.scatter_add |
| Tensor.sigmoid |
| Tensor.sign |
| Tensor.sin |
| Tensor.softmax |
| Tensor.sort |
| Tensor.sqrt |
| Tensor.squeeze |
| Tensor.sub |
| Tensor.sub_ |
| Tensor.sum |
| Tensor.to |
| Tensor.topk |
| Tensor.transpose |
| Tensor.tril |
| Tensor.triu |
| Tensor.uniform_ |
| Tensor.unsqueeze |
| Tensor.view_as |
| Tensor.where |
| torch._fused_adamw_ |
| torch._log_softmax |
| torch._softmax_backward_data |
| torch.abs |
| torch.add |
| torch.addmm |
| torch.all |
| torch.any |
| torch.arange |
| torch.bitwise_and |
| torch.bitwise_not |
| torch.bitwise_or |
| torch.bmm |
| torch.cat |
| torch.chunk |
| torch.clamp |
| torch.constant_pad_nd |
| torch.convolution |
| torch.cos |
| torch.cumsum |
| torch.div |
| torch.embedding |
| torch.eq |
| torch.exp |
| torch.fill_ |
| torch.gather |
| torch.ge |
| torch.gt |
| torch.histc |
| torch.index |
| torch.index_select |
| torch.le |
| torch.linalg.vector_norm |
| torch.log |
| torch.logical_and_ |
| torch.logical_not |
| torch.lt |
| torch.masked_fill |
| torch.matmul |
| torch.max |
| torch.maximum |
| torch.mean |
| torch.min |
| torch.minimum |
| torch.mm |
| torch.mul |
| torch.native_dropout |
| torch.neg |
| torch.nn.functional.embedding |
| torch.nn.functional.layer_norm |
| torch.nn.functional.linear |
| torch.nn.functional.mse_loss |
| torch.nn.functional.one_hot |
| torch.nn.functional.pad |
| torch.nn.functional.silu |
| torch.nn.functional.softmax |
| torch.nn.GELU |
| torch.norm |
| torch.ones |
| torch.ones_like |
| torch.pow |
| torch.reciprocal |
| torch.relu |
| torch.remainder |
| torch.repeat_interleave |
| torch.rsqrt |
| torch.scatter |
| torch.scatter_add |
| torch.sigmoid |
| torch.sign |
| torch.sin |
| torch.softmax |
| torch.sort |
| torch.split |
| torch.sqrt |
| torch.stack |
| torch.sub |
| torch.sum |
| torch.topk |
| torch.tril |
| torch.triu |
| torch.where |
| torch.zeros |
| torch_npu.npu_add_rms_norm |
| torch_npu.npu_all_gather_base_mm |
| torch_npu.npu_apply_adam_w |
| torch_npu.npu_conv2d |
| torch_npu.npu_cross_entropy_loss |
| torch_npu.npu_dtype_cast |
| torch_npu.npu_fusion_attention |
| torch_npu.npu_grouped_matmul_add_ |
| torch_npu.npu_mm_reduce_scatter_base |
| torch_npu.npu_moe_token_permute |
| torch_npu.npu_moe_token_unpermute |
| torch_npu.npu_rms_norm |
| torch_npu.npu_rotary_mul |
| torch_npu.npu_swiglu |
| Tensor.__eq__ |
| Tensor.__imul__ |
| Tensor.__isub__ |
| Tensor.__truediv__ |
| torch_npu.npu_grouped_matmul |
| Tensor.__add__ |
| Tensor.__bool__ |
| Tensor.__ge__ |
| Tensor.__gt__ |
| Tensor.__iadd__ |
| Tensor.__mul__ |
| Tensor.__or__ |
| Tensor.__radd__ |
| Tensor.__rmul__ |
| Tensor.__sub__ |
| Tensor.__pow__ |
| Tensor.__setitem__ |