fuse_linear_bn_weights

→(Tensor, Tensor)

fuse_linear_bn_weights(linear_w: Tensor, linear_b: Tensor | None, bn_rm: Tensor, bn_rv: Tensor, bn_eps: float, bn_w: Tensor | None, bn_b: Tensor | None)

source edit

Low-level form of fuse_linear_bn_eval operating on raw tensors.

Internally identical to fuse_conv_bn_weights — a Linear's weight is just a 2-D conv weight from the BN's perspective. Exposed under its own name to keep call sites self-documenting in graph-rewrite passes.

Parameters

linear_wTensor

Linear weight matrix, shape (out_features, in_features).

linear_bTensor or None

Linear bias, shape (out_features,), or None.

bn_rmTensor

BatchNorm1d running mean.

bn_rvTensor

BatchNorm1d running variance.

bn_epsfloat

BatchNorm numerical-stability epsilon.

bn_wTensor or None

BatchNorm affine

\gamma

bn_bTensor or None

BatchNorm affine

\beta

Returns

(Tensor, Tensor)

Fused (weight, bias) on the same dtype / device as linear_w.

Notes

Math:

\mathbf{W}_{\text{fused}} = \mathbf{W} \cdot \gamma / \sqrt{\sigma^2 + \epsilon}, \qquad b_{\text{fused}} = \gamma (b - \mu)/\sqrt{\sigma^2 + \epsilon} + \beta.

Examples

>>> import lucid.nn as nn
>>> from lucid.nn.utils import fuse_linear_bn_weights
>>> linear = nn.Linear(128, 64); bn = nn.BatchNorm1d(64)
>>> W, b = fuse_linear_bn_weights(
...     linear.weight, linear.bias,
...     bn.running_mean, bn.running_var, bn.eps,
...     bn.weight, bn.bias,
... )
>>> W.shape, b.shape
((64, 128), (64,))

Used by 1

lucid.nn.utils

>>> import lucid.nn as nn >>> from lucid.nn.utils import fuse_linear_bn_weights >>> linear = nn.Linear(128, 64); bn = nn.BatchNorm1d(64) >>> W, b = fuse_linear_bn_weights( ... linear.weight, linear.bias, ... bn.running_mean, bn.running_var, bn.eps, ... bn.weight, bn.bias, ... ) >>> W.shape, b.shape ((64, 128), (64,))