若API未标明支持情况，则代表该API的支持情况待验证。
API名称 |
是否支持 |
限制与说明 |
|---|---|---|
torch.nn.parameter.Parameter |
是 |
支持fp32 |
torch.nn.parameter.UninitializedParameter |
是 |
|
torch.nn.parameter.UninitializedBuffer |
是 |
|
torch.nn.Module |
是 |
支持fp32 |
torch.nn.Module.add_module |
是 |
支持fp32 |
torch.nn.Module.apply |
是 |
支持fp32 |
torch.nn.Module.bfloat16 |
是 |
支持bf16，fp32，int64 |
torch.nn.Module.buffers |
是 |
|
torch.nn.Module.children |
是 |
支持fp32 |
torch.nn.Module.cpu |
是 |
支持fp32 |
torch.nn.Module.cuda |
是 |
支持fp32 |
torch.nn.Module.double |
是 |
|
torch.nn.Module.eval |
是 |
支持fp32，int64 |
torch.nn.Module.extra_repr |
是 |
支持fp32 |
torch.nn.Module.float |
是 |
支持fp16，fp32，int64 |
torch.nn.Module.forward |
是 |
支持fp32 |
torch.nn.Module.get_buffer |
是 |
|
torch.nn.Module.get_extra_state |
是 |
|
torch.nn.Module.get_parameter |
是 |
支持fp32 |
torch.nn.Module.get_submodule |
是 |
支持fp32 |
torch.nn.Module.half |
是 |
支持fp16，fp32，int64 |
torch.nn.Module.ipu |
否 |
|
torch.nn.Module.load_state_dict |
是 |
支持fp32 |
torch.nn.Module.modules |
是 |
支持fp32 |
torch.nn.Module.named_buffers |
是 |
|
torch.nn.Module.named_children |
是 |
支持fp32 |
torch.nn.Module.named_modules |
是 |
支持fp32 |
torch.nn.Module.named_parameters |
是 |
|
torch.nn.Module.parameters |
是 |
|
torch.nn.Module.register_backward_hook |
是 |
支持fp32 |
torch.nn.Module.register_buffer |
是 |
支持fp32 |
torch.nn.Module.register_forward_hook |
是 |
支持fp32 |
torch.nn.Module.register_forward_pre_hook |
是 |
支持fp32 |
torch.nn.Module.register_full_backward_hook |
是 |
支持fp32 |
torch.nn.Module.register_full_backward_pre_hook |
否 |
|
torch.nn.Module.register_load_state_dict_post_hook |
否 |
|
torch.nn.Module.register_module |
是 |
支持fp32 |
torch.nn.Module.register_parameter |
是 |
|
torch.nn.Module.register_state_dict_pre_hook |
否 |
|
torch.nn.Module.requires_grad_ |
是 |
|
torch.nn.Module.set_extra_state |
是 |
|
torch.nn.Module.share_memory |
||
torch.nn.Module.state_dict |
是 |
支持fp32 |
torch.nn.Module.to |
是 |
支持fp32，int64 |
torch.nn.Module.to_empty |
是 |
支持fp32 |
torch.nn.Module.train |
是 |
支持fp32 |
torch.nn.Module.type |
是 |
支持fp16，fp32，int64 |
torch.nn.Module.xpu |
否 |
|
torch.nn.Module.zero_grad |
是 |
支持fp32 |
torch.nn.Sequential |
是 |
支持fp32 |
torch.nn.Sequential.append |
是 |
支持fp32 |
torch.nn.ModuleList |
是 |
支持fp32 |
torch.nn.ModuleList.append |
是 |
支持fp32 |
torch.nn.ModuleList.extend |
是 |
支持fp32 |
torch.nn.ModuleList.insert |
是 |
支持fp32 |
torch.nn.ModuleDict |
是 |
支持fp32 |
torch.nn.ModuleDict.clear |
是 |
支持fp32 |
torch.nn.ModuleDict.items |
是 |
支持fp32 |
torch.nn.ModuleDict.keys |
是 |
支持fp32 |
torch.nn.ModuleDict.pop |
是 |
支持fp32 |
torch.nn.ModuleDict.update |
是 |
支持fp32 |
torch.nn.ModuleDict.values |
是 |
支持fp32 |
torch.nn.ParameterList |
是 |
支持fp32 |
torch.nn.ParameterList.append |
是 |
支持fp32 |
torch.nn.ParameterList.extend |
是 |
支持fp32 |
torch.nn.ParameterDict |
是 |
支持fp32 |
torch.nn.ParameterDict.clear |
是 |
支持fp32 |
torch.nn.ParameterDict.copy |
是 |
支持fp32 |
torch.nn.ParameterDict.fromkeys |
是 |
支持fp32 |
torch.nn.ParameterDict.get |
是 |
支持fp32 |
torch.nn.ParameterDict.items |
是 |
支持fp32 |
torch.nn.ParameterDict.keys |
是 |
支持fp32 |
torch.nn.ParameterDict.pop |
是 |
支持fp32 |
torch.nn.ParameterDict.popitem |
是 |
支持fp32 |
torch.nn.ParameterDict.setdefault |
是 |
支持fp32 |
torch.nn.ParameterDict.update |
是 |
支持fp32 |
torch.nn.ParameterDict.values |
是 |
支持fp32 |
torch.nn.modules.module.register_module_forward_pre_hook |
是 |
支持fp32 |
torch.nn.modules.module.register_module_forward_hook |
是 |
支持fp32 |
torch.nn.modules.module.register_module_backward_hook |
是 |
支持fp32 |
torch.nn.modules.module.register_module_full_backward_hook |
是 |
支持fp32 |
torch.nn.Conv1d |
是 |
支持fp32，fp16 |
torch.nn.Conv2d |
是 |
支持fp16，fp32 |
torch.nn.Conv3d |
否 |
|
torch.nn.ConvTranspose1d |
是 |
支持fp32 |
torch.nn.ConvTranspose2d |
是 |
支持fp16，fp32 |
torch.nn.ConvTranspose3d |
否 |
|
torch.nn.LazyConv1d |
是 |
支持fp32，fp16 |
torch.nn.LazyConv2d |
是 |
支持fp32，fp16 |
torch.nn.LazyConv3d |
否 |
|
torch.nn.LazyConvTranspose1d |
是 |
支持fp16 |
torch.nn.LazyConvTranspose2d |
是 |
支持fp16，fp32 |
torch.nn.LazyConvTranspose3d |
||
torch.nn.Unfold |
||
torch.nn.Fold |
||
torch.nn.MaxPool1d |
||
torch.nn.MaxPool2d |
||
torch.nn.MaxPool3d |
||
torch.nn.MaxUnpool1d |
是 |
支持fp16，fp32 |
torch.nn.MaxUnpool2d |
是 |
支持fp16，fp32 |
torch.nn.MaxUnpool3d |
||
torch.nn.AvgPool1d |
否 |
|
torch.nn.AvgPool2d |
否 |
|
torch.nn.AvgPool3d |
否 |
|
torch.nn.FractionalMaxPool2d |
否 |
|
torch.nn.FractionalMaxPool3d |
否 |
|
torch.nn.LPPool1d |
否 |
|
torch.nn.LPPool2d |
否 |
|
torch.nn.AdaptiveMaxPool1d |
||
torch.nn.AdaptiveMaxPool2d |
||
torch.nn.AdaptiveMaxPool3d |
||
torch.nn.AdaptiveAvgPool1d |
是 |
支持fp32，fp16 |
torch.nn.AdaptiveAvgPool2d |
是 |
支持fp32，fp16 |
torch.nn.AdaptiveAvgPool3d |
||
torch.nn.ReflectionPad1d |
是 |
支持fp32，fp16 |
torch.nn.ReflectionPad2d |
是 |
支持fp32，fp16 |
torch.nn.ReflectionPad3d |
否 |
|
torch.nn.ReplicationPad1d |
是 |
支持fp32，fp16，complex64，complex128 |
torch.nn.ReplicationPad2d |
是 |
支持fp32，fp16，complex64，complex128 |
torch.nn.ReplicationPad3d |
否 |
可以走CPU实现 |
torch.nn.ZeroPad2d |
是 |
支持bf16，fp32，fp16，uint8，int8，int16，int32，int64，bool，complex64，complex128 |
torch.nn.ConstantPad1d |
是 |
支持int8，bool 在输入x为六维以上时可能会出现性能下降问题 |
torch.nn.ConstantPad2d |
是 |
支持fp32，fp16，uint8，int8，int16，int32，int64，complex64，complex128 在输入x为六维以上时可能会出现性能下降问题 |
torch.nn.ConstantPad3d |
否 |
|
torch.nn.ELU |
是 |
支持fp32，fp16 |
torch.nn.Hardshrink |
是 |
支持fp32，fp16 |
torch.nn.Hardsigmoid |
是 |
支持fp16，fp32，int32 |
torch.nn.Hardtanh |
是 |
支持fp32，fp16，int8，int16，int32，int64 |
torch.nn.Hardswish |
是 |
支持fp16，fp32 |
torch.nn.LeakyReLU |
是 |
支持fp16，fp32，fp64 |
torch.nn.LogSigmoid |
是 |
支持fp32，fp16 |
torch.nn.MultiheadAttention |
否 |
|
torch.nn.MultiheadAttention.forward |
否 |
|
torch.nn.MultiheadAttention.merge_masks |
否 |
|
torch.nn.PReLU |
是 |
支持fp32 |
torch.nn.ReLU |
是 |
支持bf16，fp16，fp32，int8，uint8，int32，int64 |
torch.nn.ReLU6 |
是 |
支持fp32，fp16，uint8，int8，int16，int32，int64 |
torch.nn.RReLU |
||
torch.nn.SELU |
是 |
支持fp16，fp32，int8，int32，fp64，uint8，int16，int64，bool |
torch.nn.CELU |
是 |
支持fp32，fp16 |
torch.nn.GELU |
是 |
支持bf16，fp32，fp16 |
torch.nn.Sigmoid |
是 |
支持bf16，fp32，fp16，uint8，int8，int16，int32，int64，bool，complex64，complex128 |
torch.nn.SiLU |
是 |
支持bf16，fp16，fp32 |
torch.nn.Mish |
||
torch.nn.Softplus |
是 |
支持fp16，fp32 |
torch.nn.Softshrink |
是 |
支持bf16，fp16，fp32 |
torch.nn.Softsign |
是 |
支持bf16，fp32，fp16，uint8，int8，int16，int32，int64 |
torch.nn.Tanh |
是 |
支持bf16，fp16，fp32，bool |
torch.nn.Tanhshrink |
是 |
支持fp32，fp16，uint8，int8，int16，int32，int64 |
torch.nn.Threshold |
是 |
支持fp16，fp32，int8，uint8，int16，int32，int64 |
torch.nn.GLU |
是 |
支持fp32，fp16 |
torch.nn.Softmin |
是 |
支持bf16，fp32，fp16 |
torch.nn.Softmax |
是 |
支持bf16，fp32，fp16 |
torch.nn.Softmax2d |
是 |
支持bf16，fp32，fp16 |
torch.nn.LogSoftmax |
是 |
支持bf16，fp32，fp16 |
torch.nn.AdaptiveLogSoftmaxWithLoss |
否 |
|
torch.nn.AdaptiveLogSoftmaxWithLoss.log_prob |
否 |
|
torch.nn.AdaptiveLogSoftmaxWithLoss.predict |
否 |
|
torch.nn.BatchNorm1d |
是 |
支持fp32，fp16 |
torch.nn.BatchNorm2d |
是 |
支持fp32，fp16 |
torch.nn.BatchNorm3d |
是 |
支持fp32，fp16 |
torch.nn.LazyBatchNorm1d |
||
torch.nn.LazyBatchNorm2d |
||
torch.nn.LazyBatchNorm3d |
||
torch.nn.GroupNorm |
是 |
支持fp32 |
torch.nn.SyncBatchNorm |
是 |
支持fp16，fp32 |
torch.nn.SyncBatchNorm.convert_sync_batchnorm |
是 |
|
torch.nn.InstanceNorm1d |
否 |
|
torch.nn.InstanceNorm2d |
否 |
|
torch.nn.InstanceNorm3d |
否 |
|
torch.nn.LazyInstanceNorm1d |
||
torch.nn.LazyInstanceNorm2d |
||
torch.nn.LazyInstanceNorm3d |
||
torch.nn.LayerNorm |
是 |
支持fp32，fp16 |
torch.nn.LocalResponseNorm |
否 |
|
torch.nn.RNNBase |
||
torch.nn.RNNBase.flatten_parameters |
||
torch.nn.RNN |
否 |
|
torch.nn.LSTM |
否 |
|
torch.nn.GRU |
否 |
|
torch.nn.RNNCell |
||
torch.nn.LSTMCell |
否 |
|
torch.nn.GRUCell |
是 |
支持fp32，fp16 |
torch.nn.Transformer |
是 |
支持fp32，fp16 |
torch.nn.Transformer.forward |
||
torch.nn.Transformer.generate_square_subsequent_mask |
否 |
|
torch.nn.TransformerEncoder |
否 |
|
torch.nn.TransformerEncoder.forward |
否 |
|
torch.nn.TransformerDecoder |
||
torch.nn.TransformerDecoder.forward |
||
torch.nn.TransformerEncoderLayer |
否 |
|
torch.nn.TransformerEncoderLayer.forward |
否 |
|
torch.nn.TransformerDecoderLayer |
否 |
|
torch.nn.TransformerDecoderLayer.forward |
否 |
|
torch.nn.Identity |
是 |
支持fp32 |
torch.nn.Linear |
是 |
支持fp32，fp16 |
torch.nn.Bilinear |
是 |
支持fp32，fp16 |
torch.nn.LazyLinear |
是 |
支持fp16，fp32 |
torch.nn.Dropout |
||
torch.nn.Dropout1d |
否 |
|
torch.nn.Dropout2d |
是 |
支持fp16，fp32，int64，bool |
torch.nn.Dropout3d |
||
torch.nn.AlphaDropout |
是 |
支持fp32，fp16 |
torch.nn.FeatureAlphaDropout |
是 |
支持fp32，fp16 |
torch.nn.Embedding |
是 |
支持int32，int64 属性max_norm仅支持非负值 |
torch.nn.Embedding.from_pretrained |
是 |
支持fp64 |
torch.nn.EmbeddingBag |
是 |
支持int32，int64 |
torch.nn.EmbeddingBag.forward |
是 |
支持int64 |
torch.nn.EmbeddingBag.from_pretrained |
是 |
支持int64 |
torch.nn.CosineSimilarity |
||
torch.nn.PairwiseDistance |
||
torch.nn.L1Loss |
是 |
支持fp32，fp16，int64， |
torch.nn.MSELoss |
是 |
支持fp32，fp16 |
torch.nn.CrossEntropyLoss |
是 |
支持fp32，fp16 |
torch.nn.CTCLoss |
是 |
支持fp32 |
torch.nn.NLLLoss |
是 |
支持fp32，fp16 |
torch.nn.PoissonNLLLoss |
是 |
支持bf16，fp32，fp16，int64 |
torch.nn.GaussianNLLLoss |
是 |
支持bf16，fp32，fp16，int16，int32，int64 |
torch.nn.KLDivLoss |
是 |
支持fp16，fp32，fp64 |
torch.nn.BCELoss |
是 |
支持fp32，fp16 |
torch.nn.BCEWithLogitsLoss |
是 |
支持fp32，fp16，int16，int32，int64，bool |
torch.nn.MarginRankingLoss |
是 |
支持bf16，fp32，fp16，int8，int16，int32，int64 |
torch.nn.HingeEmbeddingLoss |
是 |
支持bf16，fp32，fp16，uint8，int8，int16，int32，int64 |
torch.nn.MultiLabelMarginLoss |
||
torch.nn.HuberLoss |
否 |
|
torch.nn.SmoothL1Loss |
是 |
支持fp32，fp16 |
torch.nn.SoftMarginLoss |
||
torch.nn.MultiLabelSoftMarginLoss |
是 |
支持fp16，fp32 |
torch.nn.CosineEmbeddingLoss |
||
torch.nn.MultiMarginLoss |
否 |
|
torch.nn.TripletMarginLoss |
否 |
|
torch.nn.TripletMarginWithDistanceLoss |
是 |
支持bf16，fp32，fp16 |
torch.nn.PixelShuffle |
是 |
支持fp16，fp32，fp64，int8，uint8，int16，int32，int64，bool |
torch.nn.PixelUnshuffle |
是 |
支持fp16，fp32，fp64，int8，uint8，int16，int32，int64，bool |
torch.nn.Upsample |
是 |
支持fp16，fp32，fp64 只支持mode = nearest |
torch.nn.UpsamplingNearest2d |
是 |
支持fp32，fp16，uint8 |
torch.nn.UpsamplingBilinear2d |
||
torch.nn.ChannelShuffle |
是 |
支持bf16，fp32，fp16，uint8，int8，int16，int32，int64，bool，complex64，complex128 |
torch.nn.DataParallel |
否 |
|
torch.nn.parallel.DistributedDataParallel |
是 |
|
torch.nn.parallel.DistributedDataParallel.join |
否 |
|
torch.nn.parallel.DistributedDataParallel.join_hook |
否 |
|
torch.nn.parallel.DistributedDataParallel.no_sync |
是 |
|
torch.nn.parallel.DistributedDataParallel.register_comm_hook |
否 |
|
torch.nn.utils.clip_grad_norm_ |
否 |
|
torch.nn.utils.clip_grad_value_ |
是 |
支持bf16，fp32，fp16 |
torch.nn.utils.parameters_to_vector |
||
torch.nn.utils.vector_to_parameters |
是 |
支持bf16，fp16，fp32，fp64，complex64 |
torch.nn.utils.prune.BasePruningMethod |
是 |
|
torch.nn.utils.prune.BasePruningMethod.apply |
是 |
|
torch.nn.utils.prune.BasePruningMethod.apply_mask |
是 |
支持fp32 |
torch.nn.utils.prune.BasePruningMethod.prune |
是 |
支持fp32 |
torch.nn.utils.prune.BasePruningMethod.remove |
是 |
支持fp32 |
torch.nn.utils.prune.PruningContainer |
是 |
|
torch.nn.utils.prune.PruningContainer.add_pruning_method |
是 |
|
torch.nn.utils.prune.PruningContainer.apply |
||
torch.nn.utils.prune.PruningContainer.apply_mask |
是 |
|
torch.nn.utils.prune.PruningContainer.compute_mask |
是 |
支持fp32 |
torch.nn.utils.prune.PruningContainer.prune |
是 |
支持fp32 |
torch.nn.utils.prune.PruningContainer.remove |
是 |
支持fp32 |
torch.nn.utils.prune.Identity |
是 |
支持fp32 |
torch.nn.utils.prune.Identity.apply |
是 |
支持fp32 |
torch.nn.utils.prune.Identity.apply_mask |
是 |
支持fp32 |
torch.nn.utils.prune.Identity.prune |
是 |
支持fp32 |
torch.nn.utils.prune.Identity.remove |
是 |
支持fp32 |
torch.nn.utils.prune.RandomUnstructured |
是 |
支持fp32 |
torch.nn.utils.prune.RandomUnstructured.apply |
是 |
支持fp32 |
torch.nn.utils.prune.RandomUnstructured.apply_mask |
是 |
支持fp32 |
torch.nn.utils.prune.RandomUnstructured.prune |
是 |
支持fp32 |
torch.nn.utils.prune.RandomUnstructured.remove |
是 |
|
torch.nn.utils.prune.L1Unstructured |
是 |
支持fp32 |
torch.nn.utils.prune.L1Unstructured.apply |
是 |
支持fp32 |
torch.nn.utils.prune.L1Unstructured.apply_mask |
是 |
支持fp32 |
torch.nn.utils.prune.L1Unstructured.prune |
是 |
支持fp32 |
torch.nn.utils.prune.L1Unstructured.remove |
是 |
支持fp32 |
torch.nn.utils.prune.RandomStructured |
是 |
支持fp32 |
torch.nn.utils.prune.RandomStructured.apply |
是 |
支持fp32 |
torch.nn.utils.prune.RandomStructured.apply_mask |
是 |
支持fp32 |
torch.nn.utils.prune.RandomStructured.compute_mask |
是 |
支持fp32 |
torch.nn.utils.prune.RandomStructured.prune |
是 |
|
torch.nn.utils.prune.RandomStructured.remove |
是 |
|
torch.nn.utils.prune.LnStructured |
是 |
支持fp32 |
torch.nn.utils.prune.LnStructured.apply |
是 |
支持fp32 |
torch.nn.utils.prune.LnStructured.apply_mask |
是 |
支持fp32 |
torch.nn.utils.prune.LnStructured.compute_mask |
是 |
支持fp32 |
torch.nn.utils.prune.LnStructured.prune |
是 |
支持fp32 |
torch.nn.utils.prune.LnStructured.remove |
是 |
支持fp32 |
torch.nn.utils.prune.CustomFromMask |
是 |
支持int64 |
torch.nn.utils.prune.CustomFromMask.apply |
是 |
支持int64 |
torch.nn.utils.prune.CustomFromMask.apply_mask |
是 |
|
torch.nn.utils.prune.CustomFromMask.prune |
是 |
|
torch.nn.utils.prune.CustomFromMask.remove |
是 |
|
torch.nn.utils.prune.identity |
是 |
支持fp32 |
torch.nn.utils.prune.random_unstructured |
是 |
|
torch.nn.utils.prune.l1_unstructured |
是 |
|
torch.nn.utils.prune.random_structured |
是 |
|
torch.nn.utils.prune.ln_structured |
是 |
|
torch.nn.utils.prune.global_unstructured |
是 |
|
torch.nn.utils.prune.custom_from_mask |
是 |
支持int64 |
torch.nn.utils.prune.remove |
是 |
|
torch.nn.utils.prune.is_pruned |
是 |
|
torch.nn.utils.weight_norm |
||
torch.nn.utils.remove_weight_norm |
||
torch.nn.utils.spectral_norm |
是 |
|
torch.nn.utils.remove_spectral_norm |
||
torch.nn.utils.skip_init |
否 |
|
torch.nn.utils.parametrizations.orthogonal |
否 |
|
torch.nn.utils.parametrizations.spectral_norm |
否 |
|
torch.nn.utils.parametrize.register_parametrization |
否 |
|
torch.nn.utils.parametrize.remove_parametrizations |
否 |
|
torch.nn.utils.parametrize.cached |
否 |
|
torch.nn.utils.parametrize.is_parametrized |
否 |
|
torch.nn.utils.parametrize.ParametrizationList |
是 |
|
torch.nn.utils.parametrize.ParametrizationList.right_inverse |
是 |
支持fp32 |
torch.nn.utils.stateless.functional_call |
否 |
|
torch.nn.utils.rnn.PackedSequence |
是 |
支持fp32，int64 |
torch.nn.utils.rnn.PackedSequence.count |
是 |
支持fp32 |
torch.nn.utils.rnn.PackedSequence.index |
是 |
支持fp32 |
torch.nn.utils.rnn.PackedSequence.is_pinned |
是 |
|
torch.nn.utils.rnn.PackedSequence.to |
是 |
支持fp32，int64 |
torch.nn.utils.rnn.pack_padded_sequence |
||
torch.nn.utils.rnn.pad_packed_sequence |
||
torch.nn.utils.rnn.pad_sequence |
是 |
支持fp32，fp16 |
torch.nn.utils.rnn.pack_sequence |
||
torch.nn.utils.rnn.unpack_sequence |
否 |
|
torch.nn.utils.rnn.unpad_sequence |
否 |
|
torch.nn.Flatten |
是 |
支持bf16，fp32，fp16，uint8，int8，int16，int32，int64，bool，complex64，complex128 |
torch.nn.Unflatten |
是 |
支持fp16，fp32，fp64，int8，int16，uint8，int32，int64，bool |
torch.nn.modules.lazy.LazyModuleMixin |
是 |
支持fp32 |
torch.nn.modules.lazy.LazyModuleMixin.has_uninitialized_params |
是 |
支持fp32 |
torch.nn.modules.lazy.LazyModuleMixin.initialize_parameters |
否 |