modules #

Classes:

Name	Description
`BNNModule`	Base class for all Bayesian neural network modules.
`Conv1d`	Applies a 1D convolution over an input signal composed of several input planes.
`Conv2d`	Applies a 2D convolution over an input signal composed of several input planes.
`Conv3d`	Applies a 3D convolution over an input signal composed of several input planes.
`Linear`	Applies an affine transformation to the input.
`Sequential`	A sequential container for modules.

Functions:

Name	Description
`batched_forward`	Call a torch.nn.Module on inputs with arbitrary many batch dimensions rather than

BNNModule #

BNNModule(
    parametrization: (
        Parametrization | None
    ) = MaximalUpdate(),
)

Bases: Module, ABC

Base class for all Bayesian neural network modules.

Parameters:

Name	Type	Description	Default
`parametrization`	`Parametrization \| None`	The parametrization to use. Defines the initialization and learning rate scaling for the parameters of the module.	`MaximalUpdate()`

Methods:

Name	Description
`forward`	Forward pass of the module.
`parameters_and_lrs`	Get the parameters of the module and their learning rates for the chosen optimizer
`reset_parameters`	Reset the parameters of the module and set the parametrization of all children

Attributes:

Name	Type	Description
`parametrization`	`Parametrization`	Parametrization of the module.

parametrization #

parametrization: Parametrization

Parametrization of the module.

forward #

forward(
    input: Float[Tensor, "*sample batch *in_feature"],
    /,
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
    input_contains_samples: bool = False,
    parameter_samples: (
        dict[str, Float[Tensor, "*sample parameter"]] | None
    ) = None,
) -> Float[Tensor, "*sample *batch *out_feature"]

Forward pass of the module.

Parameters:

Name	Type	Description	Default
`input`	`Float[Tensor, 'sample batch in_feature']`	Input tensor.	required
`sample_shape`	`Size`	Shape of samples.	`Size([])`
`generator`	`Generator \| None`	Random number generator.	`None`
`input_contains_samples`	`bool`	Whether the input already contains samples. If True, the input is assumed to have `len(sample_shape)` many leading dimensions containing input samples (typically outputs from previous layers).	`False`
`parameter_samples`	`dict[str, Float[Tensor, '*sample parameter']] \| None`	Dictionary of parameter samples. Used to pass sampled parameters to the module. Useful to jointly sample parameters of multiple layers.	`None`

parameters_and_lrs #

parameters_and_lrs(
    lr: float, optimizer: Literal["SGD", "Adam"]
) -> list[dict[str, Tensor | float]]

Get the parameters of the module and their learning rates for the chosen optimizer and the parametrization of the module.

Parameters:

Name	Type	Description	Default
`lr`	`float`	The global learning rate.	required
`optimizer`	`Literal['SGD', 'Adam']`	The optimizer being used.	required

reset_parameters #

reset_parameters() -> None

Reset the parameters of the module and set the parametrization of all children to the parametrization of the module.

This method should be implemented by subclasses to reset the parameters of the module.

Conv1d #

Conv1d(
    in_channels: int,
    out_channels: int,
    kernel_size: _size_1_t,
    stride: _size_1_t = 1,
    padding: str | _size_1_t = 0,
    dilation: _size_1_t = 1,
    groups: int = 1,
    bias: bool = True,
    padding_mode: str = "zeros",
    layer_type: Literal[
        "input", "hidden", "output"
    ] = "hidden",
    cov: FactorizedCovariance | None = None,
    parametrization: Parametrization = MaximalUpdate(),
    device: device = None,
    dtype: dtype = None,
)

Bases: _ConvNd

Applies a 1D convolution over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size :math:(N, C_{\text{in}}, L) and output :math:(N, C_{\text{out}}, L_{\text{out}}) can be precisely described as:

.. math:: \text{out}(N_i, C_{\text{out}j}) = \text{bias}(Cj}) + \sum, k) \star \text{input}(N_i, k)}^{C_{in} - 1} \text{weight}(C_{\text{out}_j

where :math:\star is the valid cross-correlation_ operator, :math:N is a batch size, :math:C denotes a number of channels, :math:L is a length of signal sequence.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of channels in the input image.	required
`out_channels`	`int`	Number of channels produced by the convolution.	required
`kernel_size`	`_size_1_t`	Size of the convolving kernel.	required
`stride`	`_size_1_t`	Stride of the convolution.	`1`
`padding`	`str \| _size_1_t`	Padding added to both sides of the input.	`0`
`dilation`	`_size_1_t`	Spacing between kernel elements.	`1`
`groups`	`int`	Number of blocked connections from input channels to output channels.	`1`
`bias`	`bool`	If `True`, adds a learnable bias to the output.	`True`
`padding_mode`	`str`	`'zeros'`, `'reflect'`, `'replicate'` or `'circular'`. Default is `'zeros'`.	`'zeros'`
`layer_type`	`Literal['input', 'hidden', 'output']`	Type of the layer. Can be one of "input", "hidden", or "output". Controls the initialization and learning rate scaling of the parameters.	`'hidden'`
`cov`	`FactorizedCovariance \| None`	The covariance of the parameters.	`None`
`parametrization`	`Parametrization`	The parametrization to use. Defines the initialization and learning rate scaling for the parameters of the module.	`MaximalUpdate()`
`device`	`device`	The device on which to place the tensor.	`None`
`dtype`	`dtype`	The desired data type of the returned tensor.	`None`

Methods:

Name	Description
`extra_repr`
`forward`
`parameters_and_lrs`
`reset_parameters`	Reset the parameters of the module.

Attributes:

Name	Type	Description
`bias`	`Parameter`
`dilation`
`groups`
`in_channels`
`kernel_size`
`layer_type`
`out_channels`
`output_padding`
`padding`
`padding_mode`
`parametrization`	`Parametrization`	Parametrization of the module.
`params`
`stride`
`transposed`
`weight`	`Parameter`

bias #

bias: Parameter

dilation #

dilation = dilation

groups #

groups = groups

in_channels #

in_channels = in_channels

kernel_size #

kernel_size = kernel_size

layer_type #

layer_type = layer_type

out_channels #

out_channels = out_channels

output_padding #

output_padding = output_padding

padding #

padding = padding

padding_mode #

padding_mode = padding_mode

parametrization #

parametrization: Parametrization

Parametrization of the module.

params #

params = ParameterDict(mean_param_dict)

stride #

stride = stride

transposed #

transposed = transposed

weight #

weight: Parameter

extra_repr #

extra_repr()

forward #

forward(
    input: Float[Tensor, "*sample batch *in_feature"],
    /,
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
    input_contains_samples: bool = False,
    parameter_samples: (
        dict[str, Float[Tensor, "*sample parameter"]] | None
    ) = None,
) -> Float[
    Tensor, "*sample *batch out_channel *out_feature"
]

parameters_and_lrs #

parameters_and_lrs(
    lr: float, optimizer: Literal["SGD", "Adam"] = "SGD"
) -> list[dict[str, Tensor | float]]

reset_parameters #

reset_parameters() -> None

Reset the parameters of the module.

Conv2d #

Conv2d(
    in_channels: int,
    out_channels: int,
    kernel_size: _size_2_t,
    stride: _size_2_t = 1,
    padding: str | _size_2_t = 0,
    dilation: _size_2_t = 1,
    groups: int = 1,
    bias: bool = True,
    padding_mode: str = "zeros",
    layer_type: Literal[
        "input", "hidden", "output"
    ] = "hidden",
    cov: FactorizedCovariance | None = None,
    parametrization: Parametrization = MaximalUpdate(),
    device: device = None,
    dtype: dtype = None,
)

Bases: _ConvNd

Applies a 2D convolution over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size :math:(N, C_{\text{in}}, H, W) and output :math:(N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}}) can be precisely described as:

.. math:: \text{out}(N_i, C_{\text{out}j}) = \text{bias}(Cj}) + \sum(N_i, k)}^{C_{\text{in}} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input

where :math:\star is the valid 2D cross-correlation_ operator, :math:N is a batch size, :math:C denotes a number of channels, :math:H is a height of input planes in pixels, and :math:W is width in pixels.

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of channels in the input image.	required
`out_channels`	`int`	Number of channels produced by the convolution.	required
`kernel_size`	`_size_2_t`	Size of the convolving kernel.	required
`stride`	`_size_2_t`	Stride of the convolution.	`1`
`padding`	`str \| _size_2_t`	Padding added to all four sides of the input.	`0`
`dilation`	`_size_2_t`	Spacing between kernel elements.	`1`
`groups`	`int`	Number of blocked connections from input channels to output channels.	`1`
`bias`	`bool`	If `True`, adds a learnable bias to the output.	`True`
`padding_mode`	`str`	`'zeros'`, `'reflect'`, `'replicate'` or `'circular'`. Default is `'zeros'`.	`'zeros'`
`layer_type`	`Literal['input', 'hidden', 'output']`	Type of the layer. Can be one of "input", "hidden", or "output". Controls the initialization and learning rate scaling of the parameters.	`'hidden'`
`cov`	`FactorizedCovariance \| None`	The covariance of the parameters.	`None`
`parametrization`	`Parametrization`	The parametrization to use. Defines the initialization and learning rate scaling for the parameters of the module.	`MaximalUpdate()`
`device`	`device`	The device on which to place the tensor.	`None`
`dtype`	`dtype`	The desired data type of the returned tensor.	`None`

Methods:

Name	Description
`extra_repr`
`forward`
`parameters_and_lrs`
`reset_parameters`	Reset the parameters of the module.

Attributes:

Name	Type	Description
`bias`	`Parameter`
`dilation`
`groups`
`in_channels`
`kernel_size`
`layer_type`
`out_channels`
`output_padding`
`padding`
`padding_mode`
`parametrization`	`Parametrization`	Parametrization of the module.
`params`
`stride`
`transposed`
`weight`	`Parameter`

bias #

bias: Parameter

dilation #

dilation = dilation

groups #

groups = groups

in_channels #

in_channels = in_channels

kernel_size #

kernel_size = kernel_size

layer_type #

layer_type = layer_type

out_channels #

out_channels = out_channels

output_padding #

output_padding = output_padding

padding #

padding = padding

padding_mode #

padding_mode = padding_mode

parametrization #

parametrization: Parametrization

Parametrization of the module.

params #

params = ParameterDict(mean_param_dict)

stride #

stride = stride

transposed #

transposed = transposed

weight #

weight: Parameter

extra_repr #

extra_repr()

forward #

forward(
    input: Float[Tensor, "*sample batch *in_feature"],
    /,
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
    input_contains_samples: bool = False,
    parameter_samples: (
        dict[str, Float[Tensor, "*sample parameter"]] | None
    ) = None,
) -> Float[
    Tensor, "*sample *batch out_channel *out_feature"
]

parameters_and_lrs #

parameters_and_lrs(
    lr: float, optimizer: Literal["SGD", "Adam"] = "SGD"
) -> list[dict[str, Tensor | float]]

reset_parameters #

reset_parameters() -> None

Reset the parameters of the module.

Conv3d #

Conv3d(
    in_channels: int,
    out_channels: int,
    kernel_size: _size_3_t,
    stride: _size_3_t = 1,
    padding: str | _size_3_t = 0,
    dilation: _size_3_t = 1,
    groups: int = 1,
    bias: bool = True,
    padding_mode: str = "zeros",
    layer_type: Literal[
        "input", "hidden", "output"
    ] = "hidden",
    cov: FactorizedCovariance | None = None,
    parametrization: Parametrization = MaximalUpdate(),
    device: device = None,
    dtype: dtype = None,
)

Bases: _ConvNd

Applies a 3D convolution over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size :math:(N, C_{in}, D, H, W) and output :math:(N, C_{out}, D_{out}, H_{out}, W_{out}) can be precisely described as:

.. math:: out(N_i, C_{out_j}) = bias(C_{out_j}) + \sum_{k = 0}^{C_{in} - 1} weight(C_{out_j}, k) \star input(N_i, k)

where :math:\star is the valid 3D cross-correlation_ operator

Parameters:

Name	Type	Description	Default
`in_channels`	`int`	Number of channels in the input image.	required
`out_channels`	`int`	Number of channels produced by the convolution.	required
`kernel_size`	`_size_3_t`	Size of the convolving kernel.	required
`stride`	`_size_3_t`	Stride of the convolution.	`1`
`padding`	`str \| _size_3_t`	Padding added to all six sides of the input.	`0`
`dilation`	`_size_3_t`	Spacing between kernel elements.	`1`
`groups`	`int`	Number of blocked connections from input channels to output channels.	`1`
`bias`	`bool`	If `True`, adds a learnable bias to the output.	`True`
`padding_mode`	`str`	`'zeros'`, `'reflect'`, `'replicate'` or `'circular'`. Default is `'zeros'`.	`'zeros'`
`layer_type`	`Literal['input', 'hidden', 'output']`	Type of the layer. Can be one of "input", "hidden", or "output". Controls the initialization and learning rate scaling of the parameters.	`'hidden'`
`cov`	`FactorizedCovariance \| None`	The covariance of the parameters.	`None`
`parametrization`	`Parametrization`	The parametrization to use. Defines the initialization and learning rate scaling for the parameters of the module.	`MaximalUpdate()`
`device`	`device`	The device on which to place the tensor.	`None`
`dtype`	`dtype`	The desired data type of the returned tensor.	`None`

Methods:

Name	Description
`extra_repr`
`forward`
`parameters_and_lrs`
`reset_parameters`	Reset the parameters of the module.

Attributes:

Name	Type	Description
`bias`	`Parameter`
`dilation`
`groups`
`in_channels`
`kernel_size`
`layer_type`
`out_channels`
`output_padding`
`padding`
`padding_mode`
`parametrization`	`Parametrization`	Parametrization of the module.
`params`
`stride`
`transposed`
`weight`	`Parameter`

bias #

bias: Parameter

dilation #

dilation = dilation

groups #

groups = groups

in_channels #

in_channels = in_channels

kernel_size #

kernel_size = kernel_size

layer_type #

layer_type = layer_type

out_channels #

out_channels = out_channels

output_padding #

output_padding = output_padding

padding #

padding = padding

padding_mode #

padding_mode = padding_mode

parametrization #

parametrization: Parametrization

Parametrization of the module.

params #

params = ParameterDict(mean_param_dict)

stride #

stride = stride

transposed #

transposed = transposed

weight #

weight: Parameter

extra_repr #

extra_repr()

forward #

forward(
    input: Float[Tensor, "*sample batch *in_feature"],
    /,
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
    input_contains_samples: bool = False,
    parameter_samples: (
        dict[str, Float[Tensor, "*sample parameter"]] | None
    ) = None,
) -> Float[
    Tensor, "*sample *batch out_channel *out_feature"
]

parameters_and_lrs #

parameters_and_lrs(
    lr: float, optimizer: Literal["SGD", "Adam"] = "SGD"
) -> list[dict[str, Tensor | float]]

reset_parameters #

reset_parameters() -> None

Reset the parameters of the module.

Linear #

Linear(
    in_features: int,
    out_features: int,
    bias: bool = True,
    layer_type: Literal[
        "input", "hidden", "output"
    ] = "hidden",
    cov: FactorizedCovariance | None = None,
    parametrization: Parametrization = MaximalUpdate(),
    device: device | None = None,
    dtype: dtype | None = None,
)

Bases: BNNModule

Applies an affine transformation to the input.

Parameters:

Name	Type	Description	Default
`in_features`	`int`	Size of each input sample.	required
`out_features`	`int`	Size of each output sample.	required
`bias`	`bool`	If set to `False`, the layer will not learn an additive bias.	`True`
`layer_type`	`Literal['input', 'hidden', 'output']`	Type of the layer. Can be one of "input", "hidden", or "output". Controls the initialization and learning rate scaling of the parameters.	`'hidden'`
`cov`	`FactorizedCovariance \| None`	Covariance object for the parameters.	`None`
`parametrization`	`Parametrization`	The parametrization to use. Defines the initialization and learning rate scaling for the parameters of the module.	`MaximalUpdate()`
`device`	`device \| None`	Device on which to create the parameters.	`None`
`dtype`	`dtype \| None`	Data type of the parameters.	`None`

Methods:

Name	Description
`extra_repr`
`forward`
`parameters_and_lrs`
`reset_parameters`	Reset the parameters of the module.

Attributes:

Name	Type	Description
`bias`	`Parameter`
`in_features`
`layer_type`
`out_features`
`parametrization`	`Parametrization`	Parametrization of the module.
`params`
`weight`	`Parameter`

bias #

bias: Parameter

in_features #

in_features = in_features

layer_type #

layer_type = layer_type

out_features #

out_features = out_features

parametrization #

parametrization: Parametrization

Parametrization of the module.

params #

params = ParameterDict(mean_param_dict)

weight #

weight: Parameter

extra_repr #

extra_repr() -> str

forward #

forward(
    input: Float[Tensor, "*sample *batch in_feature"],
    /,
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
    input_contains_samples: bool = False,
    parameter_samples: (
        dict[str, Float[Tensor, "*sample parameter"]] | None
    ) = None,
) -> Float[Tensor, "*sample *batch out_feature"]

parameters_and_lrs #

parameters_and_lrs(
    lr: float, optimizer: Literal["SGD", "Adam"] = "SGD"
) -> list[dict[str, Tensor | float]]

reset_parameters #

reset_parameters() -> None

Reset the parameters of the module.

Sequential #

Sequential(
    *args: BNNModule,
    parametrization: Parametrization | None = None
)

Sequential(
    arg: OrderedDict[str, BNNModule],
    parametrization: Parametrization | None = None,
)

Sequential(
    *args, parametrization: Parametrization | None = None
)

Bases: BNNModule, Sequential

A sequential container for modules.

Modules will be added to it in the order they are passed in the constructor. Alternatively, an OrderedDict of modules can be passed in. The forward() method of Sequential accepts any input and forwards it to the first module it contains. It then "chains" outputs to inputs sequentially for each subsequent module, finally returning the output of the last module.

The value a Sequential provides over manually calling a sequence of modules is that it allows treating the whole container as a single module, such that performing a transformation on the Sequential applies to each of the modules it stores (which are each a registered submodule of the Sequential)

Parameters:

Name	Type	Description	Default
`*args`		Any number of modules to add to the container.	`()`
`parametrization`	`Parametrization \| None`	The parametrization to use. If `None`, the parametrization of the modules in the container will be used. If a :class:`~inferno.bnn.params.Parametrization` object is passed, it will be used for all modules in the container.	`None`

Methods:

Name	Description
`forward`
`parameters_and_lrs`	Get the parameters of the module and their learning rates for the chosen optimizer
`reset_parameters`	Reset the parameters of the module and set the parametrization of all children

Attributes:

Name	Type	Description
`parametrization`		Parametrization of the module.

parametrization #

parametrization = parametrization

Parametrization of the module.

forward #

forward(
    input: Float[Tensor, "*batch in_feature"],
    /,
    sample_shape: Size = Size([]),
    generator: Generator | None = None,
    input_contains_samples: bool = False,
    parameter_samples: (
        dict[str, Float[Tensor, "*sample parameter"]] | None
    ) = None,
) -> Float[Tensor, "*sample *batch out_feature"]

parameters_and_lrs #

parameters_and_lrs(
    lr: float, optimizer: Literal["SGD", "Adam"]
) -> list[dict[str, Tensor | float]]

Get the parameters of the module and their learning rates for the chosen optimizer and the parametrization of the module.

Parameters:

Name	Type	Description	Default
`lr`	`float`	The global learning rate.	required
`optimizer`	`Literal['SGD', 'Adam']`	The optimizer being used.	required

reset_parameters #

reset_parameters() -> None

Reset the parameters of the module and set the parametrization of all children to the parametrization of the module.

This method should be implemented by subclasses to reset the parameters of the module.

batched_forward #

batched_forward(
    obj: Module, num_batch_dims: int
) -> Callable[
    [Float[Tensor, "*sample batch *in_feature"]],
    Float[Tensor, "*sample batch *out_feature"],
]

Call a torch.nn.Module on inputs with arbitrary many batch dimensions rather than just a single one.

This is useful to extend the functionality of a torch.nn.Module to work with arbitrary many batch dimensions, for example arbitrary many sampling dimensions.

Parameters:

Name	Type	Description	Default
`obj`	`Module`	The torch.nn.Module to call.	required
`num_batch_dims`	`int`	The number of batch dimensions.	required