支持入图的API清单
昇腾支持入图的API列表请参见表1。
Note
- 标注"支持入图"的 API,表示该接口可在
torch.compile编译后的计算图中正常使用,能够被图捕获并参与优化。关于图模式的说明请参考PyTorch图模式。 - 本文档展示的列表为已验证过支持入图的API。部分未验证的API可能也具备支持入图的能力,在昇腾NPU环境中请谨慎使用。我们将持续验证更多API并定期更新本文档。
| API名称 |
|---|
| Tensor.all |
| Tensor.any |
| Tensor.argmax |
| Tensor.argmin |
| Tensor.bitwise_and |
| Tensor.bitwise_not |
| Tensor.bitwise_or |
| Tensor.bmm |
| Tensor.clamp |
| Tensor.copy_ |
| Tensor.cos |
| Tensor.cumsum |
| Tensor.div |
| Tensor.eq |
| Tensor.exponential_ |
| Tensor.fill_ |
| Tensor.gather |
| Tensor.ge |
| Tensor.gt |
| Tensor.histc |
| Tensor.index_add |
| Tensor.index_fill |
| Tensor.index_put |
| Tensor.index_select |
| 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.ne |
| Tensor.neg |
| Tensor.pow |
| Tensor.reciprocal |
| Tensor.remainder |
| Tensor.repeat |
| Tensor.repeat_interleave |
| Tensor.roll |
| Tensor.rsqrt |
| Tensor.scatter_ |
| Tensor.scatter_add |
| Tensor.sigmoid |
| Tensor.sign |
| Tensor.sin |
| Tensor.softmax |
| Tensor.sort |
| Tensor.sqrt |
| Tensor.sub |
| Tensor.sum |
| Tensor.to |
| Tensor.topk |
| Tensor.tril |
| Tensor.triu |
| Tensor.uniform_ |
| 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.argmax |
| torch.argmin |
| torch.bitwise_and |
| torch.bitwise_not |
| torch.bitwise_or |
| torch.bmm |
| torch.cat |
| torch.clamp |
| torch.constant_pad_nd |
| torch.convolution |
| torch.cos |
| torch.cumsum |
| torch.device |
| torch.distributed._all_gather_base |
| torch.distributed._reduce_scatter_base |
| torch.distributed.all_gather |
| torch.distributed.all_gather_into_tensor |
| torch.distributed.all_reduce |
| torch.distributed.all_to_all_single |
| torch.distributed.get_process_group_ranks |
| torch.distributed.get_rank |
| torch.distributed.get_world_size |
| torch.distributed.init_process_group |
| torch.distributed.is_gloo_available |
| torch.distributed.reduce_op |
| torch.distributed.reduce_scatter_tensor |
| torch.distributed.ReduceOp |
| torch.div |
| torch.dropout |
| torch.embedding |
| torch.eq |
| torch.exp |
| torch.fill_ |
| torch.floor_divide |
| torch.gather |
| torch.ge |
| torch.gt |
| torch.index |
| torch.index_add |
| torch.index_fill |
| torch.index_put |
| 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.ne |
| torch.neg |
| torch.nn.functional.cross_entropy |
| torch.nn.functional.layer_norm |
| torch.nn.functional.mse_loss |
| torch.nn.functional.one_hot |
| torch.nn.functional.scaled_dot_product_attention |
| torch.nn.functional.silu |
| torch.nn.functional.softmax |
| torch.nn.GELU |
| torch.normal |
| torch.pow |
| torch.reciprocal |
| torch.relu |
| torch.remainder |
| torch.roll |
| torch.rsqrt |
| torch.rsub |
| torch.scatter |
| torch.scatter_add |
| torch.sigmoid |
| torch.sign |
| torch.sin |
| torch.softmax |
| torch.sort |
| torch.sqrt |
| torch.stack |
| torch.sub |
| torch.sum |
| torch.threshold |
| 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_dtype_cast |
| torch_npu.npu_fusion_attention |
| 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 |
| torch_npu.npu_grouped_matmul |
| torch._foreach_copy_ |