支持入图的API清单

昇腾支持入图的API列表请参见表1

Note

  • 标注"支持入图"的 API,表示该接口可在 torch.compile 编译后的计算图中正常使用,能够被图捕获并参与优化。关于图模式的说明请参考PyTorch图模式
  • 本文档展示的列表为已验证过支持入图的API。部分未验证的API可能也具备支持入图的能力,在昇腾NPU环境中请谨慎使用。我们将持续验证更多API并定期更新本文档。

表 1 昇腾支持入图的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_