torch.nn

Convolution Layers

API名称	API状态	限制与说明
nn.Conv1d	Beta	入参不支持device；不支持circular模式；replicate模式不支持bfloat16输入；支持数据类型：fp16、fp32
nn.Conv2d	Beta	入参不支持device；不支持circular模式；replicate模式不支持bfloat16输入；支持数据类型：bf16、fp16、fp32
nn.Conv3d	Beta	入参不支持device；不支持circular模式；replicate模式不支持bfloat16输入；
nn.ConvTranspose1d	Beta	不支持out出参、入参不支持device；支持数据类型：fp32
nn.ConvTranspose2d	Beta	不支持out出参、入参不支持device；支持数据类型：fp16、fp32
nn.ConvTranspose3d	Beta	不支持out出参、入参不支持device
nn.LazyConv1d	Not Support	N/A
nn.LazyConv2d	Not Support	N/A
nn.LazyConv3d	Not Support	N/A
nn.LazyConvTranspose1d	Not Support	N/A
nn.LazyConvTranspose2d	Not Support	N/A
nn.LazyConvTranspose3d	Not Support	N/A
nn.Unfold	Beta	支持fp16
nn.Fold	Beta	支持fp16

Pooling Layers

API名称	API状态	限制与说明
nn.MaxPool1d	Beta	N/A
nn.MaxPool2d	Beta	只支持4维输入；不支持bfloat16类型输入
nn.MaxPool3d	Beta	N/A
nn.MaxUnpool1d	Not Support	N/A
nn.MaxUnpool2d	Not Support	N/A
nn.MaxUnpool3d	Not Support	N/A
nn.AvgPool1d	Stable	不支持CPU；支持数据类型：fp16、fp32
nn.AvgPool2d	Beta	支持数据类型：fp16、fp32
nn.AvgPool3d	Beta	N/A
nn.FractionalMaxPool2d	Not Support	N/A
nn.FractionalMaxPool3d	Not Support	N/A
nn.LPPool1d	Not Support	N/A
nn.LPPool2d	Not Support	N/A
nn.LPPool3d	Not Support	N/A
nn.AdaptiveMaxPool1d	Not Support	N/A
nn.AdaptiveMaxPool2d	Not Support	N/A
nn.AdaptiveMaxPool3d	Not Support	N/A
nn.AdaptiveAvgPool1d	Beta	支持数据类型：fp16、fp32
nn.AdaptiveAvgPool2d	Beta	支持数据类型：fp16、fp32
nn.AdaptiveAvgPool3d	Not Support	N/A

Padding Layers

API名称	API状态	限制与说明
nn.ReflectionPad1d	Not Support	N/A
nn.ReflectionPad2d	Not Support	N/A
nn.ReflectionPad3d	Not Support	N/A
nn.ReplicationPad1d	Not Support	N/A
nn.ReplicationPad2d	Not Support	N/A
nn.ReplicationPad3d	Not Support	N/A
nn.ZeroPad1d	Beta	支持数据类型：bf16、fp16、fp32、fp64
nn.ZeroPad2d	Beta	支持数据类型：bf16、fp16、fp32、uint8、int8、int16、int32、int64、bool
nn.ZeroPad3d	Beta	支持数据类型：bf16、fp16、fp32、fp64
nn.ConstantPad1d	Beta	支持数据类型：int8、bool
nn.ConstantPad2d	Beta	支持数据类型：fp16、fp32、uint8、int8、int16、int32、int64
nn.ConstantPad3d	Beta	支持数据类型：fp16、fp32、uint8、int8、int16、int32、int64、bool
nn.CircularPad1d	Not Support	N/A
nn.CircularPad2d	Not Support	N/A
nn.CircularPad3d	Not Support	N/A

Non-linear Activations (weighted sum, nonlinearity)

API名称	API状态	限制与说明
nn.ELU	Not Support	N/A
nn.Hardshrink	Stable	支持数据类型：fp16、fp32
nn.Hardsigmoid	Beta	支持数据类型：fp16、fp32、int32
nn.Hardtanh	Not Support	N/A
nn.Hardswish	Beta	支持数据类型：fp16、fp32
nn.LeakyReLU	Beta	支持数据类型：bf16、fp16、fp32、fp64
nn.LogSigmoid	Not Support	N/A
nn.MultiheadAttention	Beta	入参不支持device
nn.PReLU	Beta	入参不支持device；支持数据类型：fp32
nn.ReLU	Stable	支持数据类型：bf16、fp16、fp32、uint8、int8、int32、int64
nn.ReLU6	Beta	支持数据类型：fp16、fp32、uint8、int8、int16、int32、int64
nn.RReLU	Not Support	N/A
nn.SELU	Beta	支持数据类型：fp16、fp32、fp64、uint8、int8、int16、int32、int64、bool
nn.CELU	Beta	支持数据类型：fp16、fp32
nn.GELU	Stable	支持数据类型：bf16、fp16、fp32
nn.Sigmoid	Beta	支持数据类型：bf16、fp16、fp32、uint8、int8、int16、int32、int64、bool
nn.SiLU	Stable	支持数据类型：bf16、fp16、fp32
nn.Mish	Beta	支持数据类型：fp16、fp32
nn.Softplus	Stable	支持数据类型：fp16、fp32
nn.Softshrink	Not Support	N/A
nn.Softsign	Not Support	N/A
nn.Tanh	Beta	支持数据类型：bf16、fp16、fp32、bool
nn.Tanhshrink	Not Support	N/A
nn.Threshold	Not Support	N/A
nn.GLU	Beta	支持数据类型：fp16、fp32

Non-linear Activations (other)

API名称	API状态	限制与说明
nn.Softmin	Not Support	N/A
nn.Softmax	Stable	支持数据类型：bf16、fp16、fp32
nn.Softmax2d	Not Support	N/A
nn.LogSoftmax	Beta	支持数据类型：bf16、fp16、fp32
nn.AdaptiveLogSoftmaxWithLoss	Beta	入参不支持device

Normalization Layers

API名称	API状态	限制与说明
nn.BatchNorm1d	Beta	入参不支持device；支持数据类型：fp16、fp32
nn.BatchNorm2d	Beta	入参不支持device；支持数据类型：fp16、fp32
nn.BatchNorm3d	Stable	入参不支持device；支持数据类型：fp16、fp32
nn.LazyBatchNorm1d	Not Support	N/A
nn.LazyBatchNorm2d	Not Support	N/A
nn.LazyBatchNorm3d	Not Support	N/A
nn.GroupNorm	Stable	入参不支持device
nn.SyncBatchNorm	Not Support	N/A
nn.InstanceNorm1d	Beta	N/A
nn.InstanceNorm2d	Beta	N/A
nn.InstanceNorm3d	Beta	N/A
nn.LazyInstanceNorm1d	Not Support	N/A
nn.LazyInstanceNorm2d	Not Support	N/A
nn.LazyInstanceNorm3d	Not Support	N/A
nn.LayerNorm	Stable	入参不支持device；支持数据类型：bf16、fp16、fp32
nn.LocalResponseNorm	Not Support	N/A
nn.RMSNorm	Not Support	N/A

Recurrent Layers

API名称	API状态	限制与说明
nn.RNNBase	Not Support	N/A
nn.RNN	Beta	N/A
nn.LSTM	Beta	支持数据类型：fp32
nn.GRU	Beta	N/A
nn.RNNCell	Beta	入参不支持device
nn.LSTMCell	Beta	入参不支持device
nn.GRUCell	Beta	入参不支持device；支持数据类型：fp16、fp32

Transformer Layers

API名称	API状态	限制与说明
nn.Transformer	Not Support	N/A
nn.TransformerEncoder	Not Support	N/A
nn.TransformerDecoder	Not Support	N/A
nn.TransformerEncoderLayer	Not Support	N/A
nn.TransformerDecoderLayer	Not Support	N/A

Linear Layers

API名称	API状态	限制与说明
nn.Identity	Stable	入参不支持device；支持数据类型：fp32
nn.Linear	Stable	入参不支持device；支持数据类型：fp16、fp32
nn.Bilinear	Not Support	N/A
nn.LazyLinear	Not Support	N/A

Dropout Layers

API名称	API状态	限制与说明
nn.Dropout	Stable	支持数据类型：bf16、fp16、fp32
nn.Dropout1d	Beta	N/A
nn.Dropout2d	Beta	支持数据类型：fp16、fp32、int64、bool
nn.Dropout3d	Beta	N/A
nn.AlphaDropout	Beta	支持数据类型：fp16、fp32
nn.FeatureAlphaDropout	Beta	支持数据类型：fp16、fp32

Sparse Layers

API名称	API状态	限制与说明
nn.Embedding	Beta	入参不支持device；支持数据类型：int32、int64
nn.EmbeddingBag	Not Support	N/A

Distance Functions

API名称	API状态	限制与说明
nn.CosineSimilarity	Beta	N/A
nn.PairwiseDistance	Beta	N/A

Loss Functions

API名称	API状态	限制与说明
nn.L1Loss	Beta	支持数据类型：fp16、fp32、int64
nn.MSELoss	Stable	支持数据类型：fp16、fp32
nn.CrossEntropyLoss	Beta	支持数据类型：fp16、fp32
nn.CTCLoss	Beta	支持数据类型：fp32
nn.NLLLoss	Beta	支持数据类型：fp16、fp32
nn.PoissonNLLLoss	Beta	支持数据类型：bf16、fp16、fp32、int64
nn.GaussianNLLLoss	Beta	支持数据类型：bf16、fp16、fp32、int16、int32、int64
nn.KLDivLoss	Beta	支持数据类型：bf16、fp16、fp32
nn.BCELoss	Beta	支持数据类型：fp16、fp32
nn.BCEWithLogitsLoss	Beta	支持数据类型：fp16、fp32
nn.MarginRankingLoss	Beta	支持数据类型：bf16、fp16、fp32、int8、int16、int32、int64
nn.HingeEmbeddingLoss	Beta	支持数据类型：bf16、fp16、fp32、uint8、int8、int16、int32、int64
nn.MultiLabelMarginLoss	Beta	N/A
nn.HuberLoss	Beta	支持数据类型：fp32、fp64
nn.SmoothL1Loss	Beta	支持数据类型：fp16、fp32
nn.SoftMarginLoss	Beta	N/A
nn.MultiLabelSoftMarginLoss	Beta	支持数据类型：fp16、fp32
nn.CosineEmbeddingLoss	Beta	N/A
nn.MultiMarginLoss	Beta	入参不支持margin；支持数据类型：fp32、fp64
nn.TripletMarginLoss	Beta	支持数据类型：fp16、fp32、uint8、int8、int16、int32、int64
nn.TripletMarginWithDistanceLoss	Beta	支持数据类型：bf16、fp16、fp32

Vision Layers

API名称	API状态	限制与说明
nn.PixelShuffle	Beta	支持数据类型：bf16、fp16、fp32、fp64、uint8、int8、int16、int32、int64、bool
nn.PixelUnshuffle	Beta	支持数据类型：fp16、fp32、fp64、uint8、int8、int16、int32、int64、bool
nn.Upsample	Beta	支持数据类型：fp16、fp32
nn.UpsamplingNearest2d	Beta	支持数据类型：fp16、fp32、uint8
nn.UpsamplingBilinear2d	Beta	N/A

Shuffle Layers

API名称	API状态	限制与说明
nn.ChannelShuffle	Not Support	N/A

DataParallel Layers (multi-GPU, distributed)

API名称	API状态	限制与说明
nn.DataParallel	Not Support	N/A
nn.parallel.DistributedDataParallel	Not Support	N/A

Utilities

API名称	API状态	限制与说明
clip_grad_norm_	Beta	N/A
clip_grad_value_	Not Support	N/A
parameters_to_vector	Not Support	N/A
vector_to_parameters	Not Support	N/A
prune.BasePruningMethod	Not Support	N/A
prune.PruningContainer	Not Support	N/A
prune.Identity	Beta	支持数据类型：fp32
prune.RandomUnstructured	Not Support	N/A
prune.L1Unstructured	Not Support	N/A
prune.RandomStructured	Not Support	N/A
prune.LnStructured	Not Support	N/A
prune.CustomFromMask	Not Support	N/A
prune.identity	Not Support	N/A
prune.random_unstructured	Not Support	N/A
prune.l1_unstructured	Not Support	N/A
prune.random_structured	Not Support	N/A
prune.ln_structured	Not Support	N/A
prune.global_unstructured	Not Support	N/A
prune.custom_from_mask	Not Support	N/A
prune.remove	Not Support	N/A
prune.is_pruned	Not Support	N/A
weight_norm	Not Support	N/A
remove_weight_norm	Not Support	N/A
spectral_norm	Not Support	N/A
remove_spectral_norm	Not Support	N/A
skip_init	Not Support	N/A
parametrizations.orthogonal	Not Support	N/A
parametrizations.spectral_norm	Not Support	N/A
parametrize.register_parametrization	Not Support	N/A
parametrize.remove_parametrizations	Not Support	N/A
parametrize.cached	Not Support	N/A
parametrize.is_parametrized	Not Support	N/A
parametrize.ParametrizationList	Not Support	N/A
stateless.functional_call	Not Support	N/A
nn.utils.rnn.PackedSequence	Not Support	N/A
nn.utils.rnn.pack_padded_sequence	Not Support	N/A
nn.utils.rnn.pad_packed_sequence	Not Support	N/A
nn.utils.rnn.pad_sequence	Not Support	N/A
nn.utils.rnn.pack_sequence	Not Support	N/A
nn.utils.rnn.unpack_sequence	Not Support	N/A
nn.utils.rnn.unpad_sequence	Not Support	N/A
nn.Flatten	Beta	支持数据类型：bf16、fp16、fp32、uint8、int8、int16、int32、int64、bool
nn.Unflatten	Beta	支持数据类型：fp16、fp32、fp64、uint8、int8、int16、int32、int64、bool