params

Classes:

Name	Description
BNNParameter
GaussianParameter	Parameter of a BNNModule with Gaussian distribution.
FactorizedCovariance	Covariance of a Gaussian parameter with a factorized structure.
DiagonalCovariance	Covariance of a Gaussian parameter with diagonal structure.
KroneckerCovariance	Covariance of a Gaussian parameter with Kronecker structure.
LowRankCovariance	Covariance of a Gaussian parameter with low-rank structure.

BNNParameter

BNNParameter(
    hyperparameters: (
        dict[str, Float[Tensor, "*hyperparameter"]] | None
    ) = None,
)

Bases: ParameterDict, ABC

Methods:

Name	Description
sample

sample

sample(
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
) -> Float[Tensor, "*sample parameter"]

GaussianParameter

GaussianParameter(
    mean: (
        Float[Tensor, "parameter"]
        | dict[str, Float[Tensor, "parameter"]]
    ),
    cov: FactorizedCovariance,
)

Bases: BNNParameter

Parameter of a BNNModule with Gaussian distribution.

Parameters:

Name	Type	Description	Default
mean	Float[Tensor, 'parameter'] | dict[str, Float[Tensor, 'parameter']]	Mean of the Gaussian distribution.	required
cov	FactorizedCovariance	Covariance of the Gaussian distribution.	required

Methods:

Name	Description
sample

Attributes:

Name	Type	Description
cov

cov

cov = cov

sample

sample(
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
) -> (
    Float[Tensor, "*sample parameter"]
    | dict[str, Float[Tensor, "*sample parameter"]]
)

FactorizedCovariance

FactorizedCovariance(rank: int | None = None)

Bases: Module

Covariance of a Gaussian parameter with a factorized structure.

Assumes the covariance is factorized as a product of a square matrix and its transpose.

..math:: \mathbf{\Sigma} = \mathbf{S} \mathbf{S}^\top

Parameters:

Name	Type	Description	Default
rank	int | None	Rank of the covariance matrix. If None, the rank is set to the total number of mean parameters.	None

Methods:

Name	Description
factor_matmul	Multiply left factor of the covariance matrix with the input.
initialize_parameters	Initialize the covariance parameters.
reset_parameters	Reset the parameters of the covariance matrix.
to_dense	Convert the covariance matrix to a dense representation.

Attributes:

Name	Type	Description
lr_scaling	dict[str, float]	Compute the learning rate scaling for the covariance parameters.
rank

lr_scaling

lr_scaling: dict[str, float]

Compute the learning rate scaling for the covariance parameters.

rank

rank = rank

factor_matmul

factor_matmul(
    input: Float[Tensor, "*sample parameter"],
    /,
    additive_constant: (
        Float[Tensor, "*sample parameter"] | None
    ) = None,
) -> dict[str, Float[Tensor, "*sample parameter"]]

Multiply left factor of the covariance matrix with the input.

Parameters:

Name	Type	Description	Default
input	Float[Tensor, '*sample parameter']	Input tensor.	required
additive_constant	Float[Tensor, '*sample parameter'] | None	Additive constant to be added to the output.	None

initialize_parameters

initialize_parameters(
    mean_parameters: dict[str, Tensor],
) -> None

Initialize the covariance parameters.

Parameters:

Name	Type	Description	Default
mean_parameters	dict[str, Tensor]	Mean parameters of the Gaussian distribution.	required

Returns:

Type	Description
None	Covariance parameters.

reset_parameters

reset_parameters(
    mean_parameter_scales: dict[str, float] | float = 1.0,
) -> None

Reset the parameters of the covariance matrix.

Initalizes the parameters of the covariance matrix with a scale that is given by the mean parameter scales and a covariance-specific scaling that depends on the structure of the covariance matrix.

Parameters:

Name	Type	Description	Default
mean_parameter_scales	dict[str, float] | float	Scales of the mean parameters. If a dictionary keys are the names of the mean parameters. If a float, all covariance parameters are initialized with the same scale.	1.0

to_dense

to_dense() -> Float[Tensor, 'parameter parameter']

Convert the covariance matrix to a dense representation.

DiagonalCovariance

DiagonalCovariance()

Bases: FactorizedCovariance

Covariance of a Gaussian parameter with diagonal structure.

Methods:

Name	Description
factor_matmul
initialize_parameters
reset_parameters
to_dense	Convert the covariance matrix to a dense representation.

Attributes:

Name	Type	Description
lr_scaling	dict[str, float]
rank

lr_scaling

lr_scaling: dict[str, float]

rank

rank = rank

factor_matmul

factor_matmul(
    input: Float[Tensor, "*sample parameter"],
    /,
    additive_constant: (
        Float[Tensor, "*sample parameter"] | None
    ) = None,
) -> dict[str, Float[Tensor, "*sample parameter"]]

initialize_parameters

initialize_parameters(
    mean_parameters: dict[str, Tensor],
) -> None

reset_parameters

reset_parameters(
    mean_parameter_scales: dict[str, float] | float = 1.0,
) -> None

to_dense

to_dense() -> Float[Tensor, 'parameter parameter']

Convert the covariance matrix to a dense representation.

KroneckerCovariance

KroneckerCovariance(
    input_rank: int | None = None,
    output_rank: int | None = None,
)

Bases: FactorizedCovariance

Covariance of a Gaussian parameter with Kronecker structure.

Assumes the covariance is given by a Kronecker product of two matrices of size equal to the number of inputs and outputs to the layer. Each Kronecker factor is assumed to be of rank \(R \leq D\) where \(D\) is either the input or output dimension of the layer.

More precisely, the covariance is given by

\[ \begin{align*} \mathbf{\Sigma} &= \mathbf{C}_{\text{in}} \otimes \mathbf{C}_{\text{out}}\\ &= \mathbf{S}_{\text{in}}\mathbf{S}_{\text{in}}^\top \otimes \mathbf{S}_{\text{out}}\mathbf{S}_{\text{out}}^\top\\ &= (\mathbf{S}_{\text{in}} \otimes \mathbf{S}_{\text{out}}) (\mathbf{S}_{\text{in}}^\top \otimes \mathbf{S}_{\text{out}}^\top) \end{align*} \]

where \(\mathbf{S}_{\text{in}}\) and \(\mathbf{S}_{\text{out}}\) are the low-rank factors of the Kronecker factors \(\mathbf{C}_{\text{in}}\) and \(\mathbf{C}_{\text{out}}\).

Parameters:

Name	Type	Description	Default
input_rank	int | None	Rank of the input Kronecker factor. If None, assumes full rank.	None
output_rank	int | None	Rank of the output Kronecker factor. If None, assumes full rank.	None

Methods:

Name	Description
factor_matmul
initialize_parameters
reset_parameters
to_dense

Attributes:

Name	Type	Description
input_rank
lr_scaling	dict[str, float]	Compute the learning rate scaling for the covariance parameters.
output_rank
rank
sample_scale	float

input_rank

input_rank = input_rank

lr_scaling

lr_scaling: dict[str, float]

Compute the learning rate scaling for the covariance parameters.

output_rank

output_rank = output_rank

rank

rank = rank

sample_scale

sample_scale: float

factor_matmul

factor_matmul(
    input: Float[Tensor, "*sample parameter"],
    /,
    additive_constant: (
        Float[Tensor, "*sample parameter"] | None
    ) = None,
) -> dict[str, Float[Tensor, "*sample parameter"]]

initialize_parameters

initialize_parameters(
    mean_parameters: dict[str, Tensor],
) -> None

reset_parameters

reset_parameters(
    mean_parameter_scales: dict[str, float] | float = 1.0,
) -> None

to_dense

to_dense() -> Float[Tensor, 'parameter parameter']

LowRankCovariance

LowRankCovariance(rank: int)

Bases: FactorizedCovariance

Covariance of a Gaussian parameter with low-rank structure.

Assumes the covariance is factorized as a product of a matrix :math:\mathbf{S} \in \mathbb{R}^{P \times R}` and its transpose.

..math:: \mathbf{\Sigma} = \mathbf{S} \mathbf{S}^\top

Parameters:

Name	Type	Description	Default
rank	int	Rank of the covariance matrix. If None, the rank is set to the total number of mean parameters.	required

Methods:

Name	Description
factor_matmul	Multiply left factor of the covariance matrix with the input.
initialize_parameters	Initialize the covariance parameters.
reset_parameters	Reset the parameters of the covariance matrix.
to_dense	Convert the covariance matrix to a dense representation.

Attributes:

Name	Type	Description
lr_scaling	dict[str, float]	Compute the learning rate scaling for the covariance parameters.
rank

lr_scaling

lr_scaling: dict[str, float]

Compute the learning rate scaling for the covariance parameters.

rank

rank = rank

factor_matmul

factor_matmul(
    input: Float[Tensor, "*sample parameter"],
    /,
    additive_constant: (
        Float[Tensor, "*sample parameter"] | None
    ) = None,
) -> dict[str, Float[Tensor, "*sample parameter"]]

Multiply left factor of the covariance matrix with the input.

Parameters:

Name	Type	Description	Default
input	Float[Tensor, '*sample parameter']	Input tensor.	required
additive_constant	Float[Tensor, '*sample parameter'] | None	Additive constant to be added to the output.	None

initialize_parameters

initialize_parameters(
    mean_parameters: dict[str, Tensor],
) -> None

Initialize the covariance parameters.

Parameters:

Name	Type	Description	Default
mean_parameters	dict[str, Tensor]	Mean parameters of the Gaussian distribution.	required

Returns:

Type	Description
None	Covariance parameters.

reset_parameters

reset_parameters(
    mean_parameter_scales: dict[str, float] | float = 1.0,
) -> None

Reset the parameters of the covariance matrix.

Initalizes the parameters of the covariance matrix with a scale that is given by the mean parameter scales and a covariance-specific scaling that depends on the structure of the covariance matrix.

Parameters:

Name	Type	Description	Default
mean_parameter_scales	dict[str, float] | float	Scales of the mean parameters. If a dictionary keys are the names of the mean parameters. If a float, all covariance parameters are initialized with the same scale.	1.0

to_dense

to_dense() -> Float[Tensor, 'parameter parameter']

Convert the covariance matrix to a dense representation.