torch

若API未标明支持情况,则代表该API的支持情况待验证。

API名称

是否支持

限制与说明

torch.SymInt

否

  

torch.SymFloat

否

  

torch.SymBool

否

  

torch.Tag

否

  

torch.is_tensor

是

  

torch.is_storage

是

  

torch.is_complex

是

  

torch.is_conj

是

  

torch.is_floating_point

是

  

torch.is_nonzero

是

  

torch.set_default_dtype

是

  

torch.get_default_dtype

是

  

torch.set_default_device

否

  

torch.set_default_tensor_type

是

  

torch.numel

是

  

torch.set_printoptions

是

  

torch.set_flush_denormal

是

  

torch.tensor

是

  

torch.sparse_coo_tensor

否

  

torch.sparse_csr_tensor

否

  

torch.sparse_csc_tensor

否

  

torch.sparse_bsr_tensor

否

  

torch.sparse_bsc_tensor

否

  

torch.asarray

是

  

torch.as_tensor

是

  

torch.as_strided

     

torch.from_numpy

是

  

torch.from_dlpack

否

  

torch.frombuffer

是

  

torch.zeros

是

  

torch.zeros_like

是

  

torch.ones

是

  

torch.ones_like

是

  

torch.arange

是

  

torch.range

是

  

torch.linspace

是

  

torch.logspace

否

  

torch.eye

是

  

torch.empty

是

  

torch.empty_like

是

  

torch.empty_strided

是

  

torch.full

是

  

torch.full_like

是

  

torch.quantize_per_tensor

是

  

torch.quantize_per_channel

是

  

torch.dequantize

     

torch.dequantize

     

torch.complex

否

  

torch.polar

否

  

torch.heaviside

否

  

torch.adjoint

否

  

torch.argwhere

否

  

torch.cat

是

  

torch.concat

是

不支持float64,不支持8D输入。

torch.concatenate

否

  

torch.conj

否

  

torch.chunk

是

  

torch.dsplit

是

只支持float16,float32,float64

torch.column_stack

是

  

torch.dstack

是

不支持float64

torch.gather

是

  

torch.hsplit

是

  

torch.hstack

是

不支持float64

torch.index_add

是

  

torch.index_copy

否

  

torch.index_reduce

否

  

torch.index_select

否

  

torch.masked_select

是

  

torch.movedim

是

  

torch.moveaxis

是

  

torch.narrow

是

  

torch.narrow_copy

否

  

torch.nonzero

是

  

torch.permute

     

torch.reshape

是

  

torch.row_stack

是

不支持float64

torch.select

是

  

torch.scatter

是

  

torch.diagonal_scatter

否

  

torch.select_scatter

是

  

torch.slice_scatter

是

  

torch.scatter_add

是

  

torch.scatter_reduce

否

  

torch.split

     

torch.squeeze

否

  

torch.stack

是

  

torch.swapaxes

是

  

torch.swapdims

是

  

torch.t

是

  

torch.take

是

  

torch.take_along_dim

是

不支持float64,bool

torch.tensor_split

是

  

torch.tile

否

  

torch.transpose

     

torch.unbind

     

torch.unsqueeze

     

torch.vsplit

是

只支持float16,float32,float64

torch.vstack

是

  

torch.where

否

  

torch.where

是

  

torch.Generator

是

  

torch.Generator.get_state

是

  

torch.Generator.initial_seed

是

  

torch.Generator.manual_seed

是

  

torch.Generator.seed

是

  

torch.Generator.set_state

是

  

torch.seed

是

  

torch.manual_seed

是

  

torch.initial_seed

是

  

torch.get_rng_state

是

  

torch.set_rng_state

是

  

torch.bernoulli

是

  

torch.multinomial

是

  

torch.normal

是

  

torch.normal

是

  

torch.normal

是

  

torch.normal

是

  

torch.poisson

否

可以走CPU实现

torch.rand

否

  

torch.rand_like

是

不支持int64

torch.randint

是

  

torch.randint_like

是

  

torch.randn

否

  

torch.randn_like

是

  

torch.randperm

是

  

torch.quasirandom.SobolEngine

是

  

torch.quasirandom.SobolEngine.draw

     

torch.quasirandom.SobolEngine.draw_base2

     

torch.quasirandom.SobolEngine.fast_forward

     

torch.quasirandom.SobolEngine.reset

     

torch.save

是

  

torch.load

     

torch.get_num_threads

是

只支持CPU,GPU/NPU不支持

torch.set_num_threads

是

只支持CPU,GPU/NPU不支持

torch.get_num_interop_threads

是

  

torch.set_num_interop_threads

是

只支持CPU,GPU/NPU不支持

torch.no_grad

是

  

torch.enable_grad

是

  

torch.set_grad_enabled

是

  

torch.is_grad_enabled

是

  

torch.inference_mode

否

  

torch.is_inference_mode_enabled

是

  

torch.abs

是

  

torch.absolute

是

  

torch.acos

是

不支持int32

torch.arccos

是

不支持int64

torch.acosh

是

不支持int64

torch.arccosh

是

不支持int64

torch.add

否

  

torch.addcdiv

是

支持fp16,fp32,int64,bool

在int64类型不支持三个tensor同时广播

torch.addcmul

是

支持fp16,fp32,fp64,int8,int32,int64,uint8,bool,bf16

在int8、uint8、int64、fp64类型不支持三个tensor同时广播

torch.angle

     

torch.asin

是

  

torch.arcsin

是

不支持int64

torch.asinh

是

  

torch.arcsinh

是

不支持int64

torch.atan

是

  

torch.arctan

是

不支持int64

torch.atanh

是

不支持int64

torch.arctanh

是

  

torch.atan2

是

  

torch.arctan2

是

  

torch.bitwise_not

是

  

torch.bitwise_and

是

  

torch.bitwise_or

是

  

torch.bitwise_xor

是

  

torch.bitwise_left_shift

否

  

torch.bitwise_right_shift

否

  

torch.ceil

是

  

torch.clamp

是

  

torch.clip

是

  

torch.conj_physical

     

torch.copysign

是

  

torch.cos

是

  

torch.cosh

是

  

torch.deg2rad

是

  

torch.div

是

  

torch.divide

否

  

torch.digamma

否

可以走CPU实现

torch.erf

是

  

torch.erfc

是

  

torch.erfinv

是

  

torch.exp

是

  

torch.exp2

     

torch.expm1

是

  

torch.fake_quantize_per_channel_affine

否

  

torch.fake_quantize_per_tensor_affine

否

  

torch.fix

否

  

torch.float_power

是

不支持float64

torch.floor

是

  

torch.floor_divide

否

  

torch.fmod

否

  

torch.frac

是

  

torch.frexp

否

  

torch.gradient

是

不支持8D输入,不支持bool,float64

torch.imag

否

  

torch.ldexp

是

在int64带out场景下,out场景也必须是int64类型。

torch.lerp

是

  

torch.lgamma

否

可以走CPU实现

torch.log

是

  

torch.log10

是

不支持bool, int8, int16,int32,int64,float64

torch.log1p

是

  

torch.log2

是

不支持bool, int8, int16,int32,int64,float64

torch.logaddexp

否

不支持double数据类型。

torch.logaddexp2

否

不支持double数据类型。

torch.logical_and

是

  

torch.logical_not

是

  

torch.logical_or

是

只支持bool

torch.logical_xor

     

torch.logit

否

可以走CPU实现

torch.hypot

否

  

torch.i0

否

  

torch.igamma

否

  

torch.igammac

否

  

torch.mul

是

  

torch.multiply

是

不支持bool

torch.mvlgamma

否

可以走CPU实现

torch.nan_to_num

否

  

torch.neg

是

  

torch.negative

是

  

torch.nextafter

否

  

torch.polygamma

否

  

torch.positive

是

  

torch.pow(input, exponent, *, out=None) ->Tensor

是

不支持int64

Int类型使用的是浮点数方案,可能存在精度误差

torch.pow(self, exponent, *, out=None) ->Tensor

是

不支持int64

Int类型使用的是浮点数方案,可能存在精度误差

torch.quantized_batch_norm

否

  

torch.quantized_max_pool1d

否

  

torch.quantized_max_pool2d

否

  

torch.rad2deg

是

不支持int64

torch.real

是

  

torch.reciprocal

是

不支持uint8,int8,int16,int32,int64,bool

torch.remainder

是

不支持int8,unit8,bool

torch.round

是

  

torch.rsqrt

是

不支持uint8,int8,int16,int32,int64,bool

torch.sigmoid

是

  

torch.sign

是

  

torch.sgn

否

  

torch.signbit

否

  

torch.sin

是

不支持uint8,int8,int16,int32,int64,bool

torch.sinc

否

  

torch.sinh

是

不支持uint8,int8,int16,int32,int64,bool

torch.softmax

否

  

torch.sqrt

是

不支持uint8,int8,int16,int32,int64,bool

torch.square

是

bool类型不支持out参数

torch.sub

是

不支持uint8

torch.subtract

是

不支持bool

torch.tan

是

不支持int8,int16,int64,bool

torch.tanh

是

  

torch.true_divide

是

不支持out参数

torch.trunc

是

只支持float16,float32

torch.xlogy

否

  

torch.argmax

否

  

torch.argmax

否

  

torch.argmin

否

  

torch.amax

是

  

torch.amin

是

  

torch.aminmax

否

  

torch.all

是

  

torch.all

是

  

torch.any

是

  

torch.any

是

  

torch.max

是

  

torch.max

是

  

torch.max

是

  

torch.min

是

  

torch.min

是

  

torch.min

是

  

torch.dist

是

  

torch.logsumexp

是

  

torch.mean

是

  

torch.mean

是

  

torch.nanmean

否

  

torch.median

是

只支持float16,float32

torch.median

是

只支持float16,float32

torch.nanmedian

否

可以走CPU实现

torch.nanmedian

否

可以走CPU实现

torch.mode

否

可以走CPU实现

torch.norm

是

  

torch.nansum

否

可以走CPU实现

torch.nansum

否

可以走CPU实现

torch.prod

是

  

torch.prod

是

  

torch.quantile

是

  

torch.nanquantile

否

  

torch.std

否

  

torch.std_mean

否

  

torch.sum

否

  

torch.sum

否

  

torch.unique

是

  

torch.unique_consecutive

是

传参时必须使用关键字,否则精度不达标。return_inverse=return_inverse,return_counts=return_counts,dim=dim。

torch.var

否

  

torch.var_mean

否

  

torch.count_nonzero

     

torch.allclose

是

  

torch.argsort

否

  

torch.eq

否

  

torch.equal

是

  

torch.ge

是

  

torch.greater_equal

是

  

torch.gt

否

  

torch.greater

是

  

torch.isclose

是

  

torch.isfinite

是

  

torch.isin

否

  

torch.isinf

否

  

torch.isposinf

是

  

torch.isneginf

是

  

torch.isnan

Atlas 训练系列产品:是

Atlas A2 训练系列产品:否

  

torch.isreal

是

  

torch.kthvalue

否

  

torch.le

否

  

torch.less_equal

是

  

torch.lt

否

  

torch.less

是

  

torch.maximum

是

  

torch.minimum

是

  

torch.fmax

否

  

torch.fmin

     

torch.ne

否

  

torch.not_equal

是

  

torch.sort

否

  

torch.topk

否

  

torch.msort

是

不支持float64,bool

torch.stft

否

  

torch.istft

否

  

torch.bartlett_window

否

  

torch.blackman_window

否

  

torch.hamming_window

否

  

torch.hann_window

否

  

torch.kaiser_window

否

  

torch.atleast_1d

是

  

torch.atleast_2d

是

  

torch.atleast_3d

是

  

torch.bincount

是

  

torch.block_diag

是

  

torch.broadcast_tensors

否

  

torch.broadcast_to

是

  

torch.broadcast_shapes

是

  

torch.bucketize

否

可以走CPU实现

torch.cartesian_prod

是

  

torch.cdist

是

  

torch.clone

是

  

torch.combinations

是

r不能大于8

torch.corrcoef

是

不支持float64

torch.cov

是

不支持bool,float16,float64

torch.cross

否

两个输入的shape要保持一致

torch.cummax

是

  

torch.cummin

是

  

torch.cumprod

是

  

torch.cumsum

否

  

torch.diag

是

仅支持diagonal=0场景

torch.diag_embed

是

不支持复数

torch.diagflat

是

  

torch.diagonal

是

仅支持diagonal=0场景

torch.diff

是

  

torch.einsum

是

  

torch.flatten

是

  

torch.flip

是

  

torch.fliplr

是

  

torch.flipud

是

  

torch.kron

是

不支持5维度及以上输入

torch.rot90

否

  

torch.gcd

     

torch.histc

否

  

torch.histogram

否

  

torch.histogramdd

否

  

torch.meshgrid

是

  

torch.lcm

否

  

torch.logcumsumexp

否

  

torch.ravel

是

  

torch.renorm

是

支持fp16,fp32

属性max_norm仅支持非负值

torch.repeat_interleave

是

  

torch.repeat_interleave

是

  

torch.roll

是

  

torch.searchsorted

是

不支持float64,bool

torch.tensordot

是

  

torch.trace

是

支持fp16,fp32,fp64,int8,int16,int32,int64,uint8,complex64,complex128

torch.tril

是

  

torch.tril_indices

是

  

torch.triu

是

  

torch.triu_indices

是

  

torch.unflatten

是

  

torch.vander

否

  

torch.view_as_real

否

  

torch.view_as_complex

否

  

torch.resolve_conj

否

  

torch.resolve_neg

否

  

torch.addbmm

否

  

torch.addmm

是

  

torch.addmv

是

  

torch.addr

是

不支持int32

torch.baddbmm

是

  

torch.bmm

是

  

torch.chain_matmul

是

  

torch.cholesky

否

  

torch.cholesky_inverse

否

  

torch.cholesky_solve

否

  

torch.dot

是

  

torch.geqrf

否

  

torch.ger

是

  

torch.inner

是

只支持float16,float32

torch.inverse

否

  

torch.det

     

torch.logdet

否

  

torch.slogdet

是

  

torch.lu

否

  

torch.lu_solve

否

  

torch.lu_unpack

否

  

torch.matmul

是

  

torch.matrix_power

否

  

torch.matrix_exp

否

  

torch.mm

是

只支持float16,float32

torch.mv

是

  

torch.orgqr

否

  

torch.ormqr

否

  

torch.outer

是

  

torch.pinverse

否

  

torch.qr

是

  

torch.svd

否

  

torch.svd_lowrank

否

可以走CPU实现

torch.pca_lowrank

否

可以走CPU实现

torch.lobpcg

否

  

torch.trapz

是

  

torch.trapezoid

否

  

torch.cumulative_trapezoid

否

  

torch.triangular_solve

否

  

torch.vdot

否

可以走CPU实现

torch._foreach_abs

否

  

torch._foreach_abs_

否

  

torch._foreach_acos

否

  

torch._foreach_acos_

否

  

torch._foreach_asin

否

  

torch._foreach_asin_

否

  

torch._foreach_atan

否

  

torch._foreach_atan_

否

  

torch._foreach_ceil

否

  

torch._foreach_ceil_

否

  

torch._foreach_cos

否

  

torch._foreach_cos_

否

  

torch._foreach_cosh

否

  

torch._foreach_cosh_

否

  

torch._foreach_erf

否

  

torch._foreach_erf_

否

  

torch._foreach_erfc

否

  

torch._foreach_erfc_

否

  

torch._foreach_exp

否

  

torch._foreach_exp_

否

  

torch._foreach_expm1

否

  

torch._foreach_expm1_

否

  

torch._foreach_floor

否

  

torch._foreach_floor_

否

  

torch._foreach_log

否

  

torch._foreach_log_

否

  

torch._foreach_log10

否

  

torch._foreach_log10_

否

  

torch._foreach_log1p

否

  

torch._foreach_log1p_

否

  

torch._foreach_log2

否

  

torch._foreach_log2_

否

  

torch._foreach_neg

否

  

torch._foreach_neg_

否

  

torch._foreach_tan

否

  

torch._foreach_tan_

否

  

torch._foreach_sin

否

  

torch._foreach_sin_

否

  

torch._foreach_sinh

否

  

torch._foreach_sinh_

否

  

torch._foreach_round

否

  

torch._foreach_round_

否

  

torch._foreach_sqrt

否

  

torch._foreach_sqrt_

否

  

torch._foreach_lgamma

否

  

torch._foreach_lgamma_

否

  

torch._foreach_frac

否

  

torch._foreach_frac_

否

  

torch._foreach_reciprocal

否

  

torch._foreach_reciprocal_

否

  

torch._foreach_sigmoid

否

  

torch._foreach_sigmoid_

否

  

torch._foreach_trunc

否

  

torch._foreach_trunc_

否

  

torch._foreach_zero_

否

  

torch.compiled_with_cxx11_abi

是

  

torch.result_type

是

  

torch.can_cast

是

  

torch.promote_types

是

  

torch.use_deterministic_algorithms

是

  

torch.are_deterministic_algorithms_enabled

是

  

torch.is_deterministic_algorithms_warn_only_enabled

否

  

torch.set_deterministic_debug_mode

否

  

torch.get_deterministic_debug_mode

否

  

torch.set_float32_matmul_precision

否

  

torch.get_float32_matmul_precision

否

  

torch.set_warn_always

否

  

torch.is_warn_always_enabled

否

  

torch.vmap

否

  

torch._assert

是

  

torch.sym_float

否

  

torch.sym_int

否

  

torch.sym_max

否

  

torch.sym_min

否

  

torch.sym_not

否

  

torch.compile

否

  

torch.autograd.set_multithreading_enabled

否