scvi.external.drvi.StackedLinearLayer

class scvi.external.drvi.StackedLinearLayer(n_stacks, in_features, out_features, bias=True, device=None, dtype=None)

Bases: Module

A parallel stacked linear layer that applies multiple linear transformations in parallel.

This layer applies a linear transformation to multiple stacks/splits of the input. It’s particularly useful in additive decoders where different splits should be calculated in parallel.

Parameters:

• n_stacks (int) – Number of stacks/splits to process in parallel.

• in_features (int) – Number of input features per stack.

• out_features (int) – Number of output features per stack.

• bias (bool (default: True)) – Whether to include bias terms for each stack.

• device (Any (default: None)) – Device to place the layer on.

• dtype (Any (default: None)) – Data type for the layer parameters.

Notes

The layer maintains separate weight and bias parameters for each stack: - Weight shape: (n_stacks, out_features, in_features) - Bias shape: (n_stacks, out_features) if bias=True, None otherwise

The forward pass applies the transformation to each stack independently: output[b, s, o] = sum_i(x[b, s, i] * weight[s, o, i]) + bias[s, o]

This is equivalent to applying n_stacks separate linear layers in parallel, which is more efficient than using separate nn.Linear layers.

Examples

>>> import torch
>>> # Create a stacked linear layer with 4 stacks
>>> layer = StackedLinearLayer(n_stacks=4, in_features=64, out_features=128)
>>> # Input shape: (batch_size, n_stacks, in_features)
>>> x = torch.randn(32, 4, 64)
>>> # Forward pass
>>> output = layer(x)
>>> print(output.shape)  # torch.Size([32, 4, 128])
>>> # Each stack has its own parameters
>>> print(layer.weight.shape)  # torch.Size([4, 128, 64])
>>> print(layer.bias.shape)  # torch.Size([4, 128])

Attributes table

n_stacks
in_features
out_features
weight
bias
training

Methods table

extra_repr()	String representation for printing the layer.
forward(x[, output_subset, stack_subset])	Forward pass through the stacked linear layer.
reset_parameters()	Reset the layer parameters to their initial values.

Attributes

StackedLinearLayer.n_stacks: int

StackedLinearLayer.in_features: int

StackedLinearLayer.out_features: int

StackedLinearLayer.weight: torch.Tensor

StackedLinearLayer.bias: torch.Tensor | None

StackedLinearLayer.training: bool

Methods

StackedLinearLayer.extra_repr()

String representation for printing the layer.

Return type:

str

StackedLinearLayer.forward(x, output_subset=None, stack_subset=None)

Forward pass through the stacked linear layer.

Parameters:

• x (Tensor) – Input tensor with shape (…, n_stacks, in_features). Any number of leading (batch) dimensions is supported; the transformation is only applied to the last two dimensions (stacks and features).

• output_subset (Tensor | None (default: None)) – Subset of outputs to provide in the output.

• stack_subset (Tensor | None (default: None)) – Indices for stacks in operation.

Return type:

Tensor

Returns:

torch.Tensor Output tensor with shape (…, n_stacks, out_features), matching the leading dimensions of the input.

Notes

The forward pass applies the linear transformation to each stack:

where: - …: arbitrary leading (batch) dimensions - s: stack index - i: input feature index - o: output feature index

The computation is performed efficiently using torch.bmm or broadcasting.

Examples

>>> import torch
>>> # Create layer
>>> layer = StackedLinearLayer(n_stacks=3, in_features=10, out_features=5)
>>> # Input: batch_size=2, n_stacks=3, in_features=10
>>> x = torch.randn(2, 3, 10)
>>> # Forward pass
>>> output = layer(x)
>>> print(output.shape)  # torch.Size([2, 3, 5])
>>> # Extra leading dimensions are also supported
>>> x = torch.randn(7, 2, 3, 10)
>>> layer(x).shape  # torch.Size([7, 2, 3, 5])

StackedLinearLayer.reset_parameters()

Reset the layer parameters to their initial values.

This method reinitializes both weights and biases using the same initialization strategy as the default nn.Linear layer.

Return type:

None

Notes

The initialization follows PyTorch’s default linear layer initialization: - Weights: Uniform distribution in [-1/sqrt(in_features), 1/sqrt(in_features)] - Biases: Uniform distribution in [-1/sqrt(in_features), 1/sqrt(in_features)]

This ensures that the variance of the output is approximately preserved across the layer.