PackedLayerNorm¶

class torch_uncertainty.layers.PackedLayerNorm(embed_dim, num_estimators, alpha, eps=1e-05, affine=True, 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, optional) – a value added to the denominator for numerical stability. Defaults to 1e-5.
affine (bool, optional) – 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.
device (torch.device, optional) – the device to use for the layer’s parameters. Defaults to None.
dtype (torch.dtype, optional) – the dtype to use for the layer’s parameters. Defaults to None.

Shape:

Input: \((B, *)\) where \(*\) means any number of additional dimensions.
Output: \((B, *)\) (same shape as input)