lazymatrix.torch.NormalizedLinear

LazyMatrix for Python · 0.2.0

Source

Bases: torch.nn.Linear

NormalizedLinear(
    in_features: int,
    out_features: int,
    bias: bool = True,
    *,
    centers: typing.Any = None,
    scales: typing.Any = None,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None
)

Apply a linear layer to lazily centered and scaled input features.

Parameters

in_features : int

Number of input features.

out_features : int

Number of output features.

bias : bool, default=True

Learn an additive bias in normalized coordinates. Disabling bias does not disable the centering correction.

centers : real array-like or torch.Tensor, optional

Fixed column centers, with one entry per input feature. None disables centering. Existing LazyMatrix centers can be passed directly.

scales : real array-like or torch.Tensor, optional

Fixed column divisors, with one entry per input feature. None disables scaling. Exact zeros, including negative zero and values that round to zero in the layer's dtype, are rejected. Negative and nonfinite values are retained. Computed LazyMatrix scales already replace zero with one.

device : torch.device or str, optional

CPU or CUDA device, following PyTorch's default device when omitted.

dtype : torch.dtype, optional

torch.float32 or torch.float64, following PyTorch's default floating dtype when omitted. Supplied statistics do not determine this default.

Notes

The layer accepts dense tensors with shape (*, in_features) and two-dimensional sparse COO or CSR tensors. Inputs must match the layer's dtype and device. Sparse inputs must not require gradients. COO duplicates are summed; CSR inputs must satisfy PyTorch's storage invariants. Other sparse layouts require explicit conversion to COO or CSR.

Weights have shape (out_features, in_features) and remain trainable in normalized coordinates. Each forward pass derives A = weight / scales and evaluates X @ A.T - A @ centers + bias, omitting disabled components. The layer never constructs the normalized input. All forward and backward operations use PyTorch, with no Rust execution or custom autograd kernels.

Centers and scales are copied, detached, persistent buffers. They move with the layer and are included in its state_dict. Checkpoint reconstruction must enable the same normalization components; saved centers and scales can be supplied to the constructor before loading the state_dict. Statistics are never fitted or updated by this layer, including in training mode.

Factoring normalization changes floating-point evaluation order. Large offsets can lose centered variation, and sparse implicit zeros and nonfinite values can differ from explicit dense normalization. Mixed precision, autocast, and devices other than CPU and CUDA are unsupported.

Examples

>>> import torch
>>> from lazymatrix.torch import NormalizedLinear
>>> layer = NormalizedLinear(2, 3, centers=[1.0, 2.0], scales=[2.0, 4.0])
>>> layer(torch.tensor([[1.0, 0.0]]).to_sparse_csr()).shape
torch.Size([1, 3])

centers

centers: Tensor | None

Source

scales

scales: Tensor | None

Source

__init__ {#init}

__init__(
    in_features: int,
    out_features: int,
    bias: bool = True,
    *,
    centers: typing.Any = None,
    scales: typing.Any = None,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None
) -> None

Source

forward

forward(input: torch.Tensor) -> torch.Tensor

Source

Transform raw features without constructing their normalized matrix.

Returns a dense tensor with the input's leading dimensions and out_features as the last dimension. Weight and bias gradients use normalized coordinates; dense input gradients use raw coordinates.

extra_repr

extra_repr() -> str

Source

Describe the layer dimensions and enabled normalization components.