torch

torch layer implementations.

Classes

BatchEnsembleConv2d

BatchEnsemble convolutional layer based on [WTB20].

BatchEnsembleLinear

BatchEnsemble linear layer based on [WTB20].

BayesConv2d

Implementation of a Bayesian convolutional layer based on [BCKW15].

BayesLinear

Implement a Bayesian linear layer based on [BCKW15].

DropConnectLinear

Custom Linear layer with DropConnect applied to weights during training.

GVBLLLayer

Generative variational Bayesian last layer based on [HWS24].

GaussianMixtureHead

Per-class Gaussian density head (Gaussian Discriminant Analysis).

HetVBLLLayer

Heteroscedastic variational Bayesian last layer based on [HWS24].

HeteroscedasticLayer

A unified PyTorch implementation of the Heteroscedastic layer.

IRDHead

Head that converts encoded features into Dirichlet concentration parameters (alpha).

IntBatchNorm1d

Interval-valued batch normalization for 1D features based on [WCM+24].

IntBatchNorm2d

Interval-valued batch normalization for 2D feature maps based on [WCM+24].

IntConv2d

Interval-arithmetic 2D convolution based on [WCM+24].

IntLinear

Interval-arithmetic linear layer based on [WCM+24].

IntSoftmax

Interval SoftMax head based on [WCM+24].

MahalanobisHead

Class-conditional Gaussian head with a tied covariance (Mahalanobis OOD).

Masksembles2D

Masksembles mask layer for 2D convolutional feature maps.

MasksemblesLinear

Masksembles mask layer for linear (dense) layers.

NatPNClassHead

Dirichlet posterior head for evidential classification.

NatPNRegHead

Gaussian posterior head for evidential regression.

NormalInverseGammaLinear

Custom Linear layer for a normal-inverse-gamma-distribution based on [ASSR20].

RadialNormalizingFlow

Radial normalizing flow based on [RM15].

RadialNormalizingFlowStack

Stack of radial normalizing flows based on [RM15].

RegressionHead

Head that converts encoded features into evidential Normal-Gamma parameters.

SNCoeffParametrization

Weight parametrization that clips the spectral norm to at most coeff.

SNGPLayer

Spectral-normalized Neural Gaussian Process (SNGP) layer based on [LLP+20].

SharedMaskDropout

Dropout that draws a single mask per forward pass, shared across the batch.

TVBLLLayer

Student-t variational Bayesian last layer based on [HWS24].

VBLLLayer

Variational Bayesian last layer based on [HWS24].

Functions

apply_spectral_norm_to_encoder

Apply spectral normalization in-place to all Conv2d and Linear layers.

pack_interval

Promote x to a packed interval tensor with lo == hi == x.

unpack_interval

Split a packed interval tensor on channel_dim into its (lo, hi) halves.