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 :math:R \leq D where :math: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 :math:\mathbf{S}_{\text{in}} and :math:\mathbf{S}_{\text{out}} are the low-rank factors of the Kronecker factors :math:\mathbf{C}_{\text{in}} and :math:\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.