PackedLayerNorm#

class torch_uncertainty.layers.PackedLayerNorm(embed_dim, num_estimators, alpha, eps=1e-05, affine=True, first=False, last=False, device=None, dtype=None)[source]#

Packed-Ensembles-style LayerNorm layer.

Parameters:

  • embed_dim (int) – the number of features in the input tensor.

  • num_estimators (int) – the number of estimators in the ensemble.

  • alpha (float) – the width multiplier of the layer.

  • eps (float) – a value added to the denominator for numerical stability. Defaults to 1e-5.

  • affine (bool) – a boolean value that when set to True, this module has learnable per_channel affine parameters initialized to ones (for weights) and zeros (for biases). Defaults to True.

  • first (bool) – Whether this layer processes the raw inputs of the model. Defaults to False.

  • last (bool) – Whether this layer processes the final outputs of the model. Defaults to False.

  • device – The device to use for the layer’s parameters. Defaults to None.

  • dtype – The dtype to use for the layer’s parameters. Defaults to None.

Shape:

  • Input: \((N, *)\) where \(*\) means any number of additional dimensions.

  • Output: \((N, *)\) (same shape as input)

Warning

This layer is only suitable to replace nn.LayerNorm when only the last dimension is normalized.