Modules

SONNModule is the base of the model, its layers and its neurons. SigmaSquashNorm is the squash of the orthogonal families (see Orthogonal polynomials), and SoftBinner the default multi-class head (see Classification).

SONNModule

SONNModule(*args: Any, **kwargs: Any)

Bases: Module

nn.Module that saves plain-Python metadata next to its tensors.

Subclasses list attribute names in params_metadata_names. state_dict stores their values under a params_metadata key, and load_state_dict restores them on every SONNModule in the tree, so settings that are not tensors travel with the checkpoint.

state_dict

state_dict(*args: Any, **kwargs: Any) -> dict[str, Any]

Return the tensor state dict plus a params_metadata entry holding the attributes named in params_metadata_names.

load_state_dict

load_state_dict(state_dict: Mapping[str, Any], *args: Any, **kwargs: Any) -> SONNModule

Load parameters and metadata recursively.

SigmaSquashNorm

SigmaSquashNorm(mean: float | Sequence[float] | Tensor, std: float | Sequence[float] | Tensor, n_sigma: float = 2.0, core_range: float = 0.75, eps: float = 1e-08)

Bases: SONNModule

Standardize by mean/std, then squash onto (-1, 1) with a rational tail.

Parameters:

Name	Type	Description	Default
mean	float, sequence of float or torch.Tensor	Per-feature centers. Scalar, or any shape broadcastable against the trailing axes of the input (typically [num_features]).	required
std	float, sequence of float or torch.Tensor	Per-feature scales, same shape as mean. Must be non-negative; entries at or below eps are treated as constant features and given a unit scale, so they map to a constant 0.	required
n_sigma	float	Half-width of the linear core, in standard deviations. Inputs within mean +/- n_sigma * std are mapped linearly onto +/-core_range; everything beyond is squashed by the rational tail. Larger values keep more of the distribution linear at the cost of resolution in the core.	2.0
core_range	float	Output magnitude reached at exactly n_sigma. Must lie in (0, 1) — at 1 the tail would have no room left to saturate into.	0.75
eps	float	Threshold below which a std entry counts as a constant feature.	1e-08

set_stats

set_stats(mean: float | Sequence[float] | Tensor, std: float | Sequence[float] | Tensor) -> None

Refit the statistics in place, keeping the buffers' shape and device.

Used to calibrate an already-constructed layer once the data that will flow through it is known — see BaseOrthogonalNeuron.fit_input_stats, which calls this per layer during training. Copies in place rather than rebinding so any vmap/functional_call view of the buffers stays valid.

from_data classmethod

from_data(x: Tensor, dim: int = 0, **kwargs: float) -> SigmaSquashNorm

Build a layer from the per-feature statistics of x.

Fit on the training split only — reusing it for validation/inference is the point, since the squash must be the same map every time.

forward

forward(x: Tensor) -> torch.Tensor

Map x into (-1, 1): linear core, saturating rational tail.

Parameters:

Name	Type	Description	Default
x	tensor	Input, broadcast against the mean and std buffers.	required

Returns:

Name	Type	Description
y	tensor	With s = (x - mean) / (n_sigmastd): core_ranges where |s| <= 1, and sign(s)*(1 - gap(|s| - 1)) beyond, with the gap of the module docstring. Same shape as the broadcast input.

extra_repr

extra_repr() -> str

Show n_sigma, core_range and the shape of the statistics in the module repr.

SoftBinner

SoftBinner(n_bins: int = 10, scale: float = 100.0)

Bases: Module

Soft one-hot encoding of a scalar prediction over n_bins fixed centers.

Produces logits (pre-softmax) — the SONN multi-class loss path applies log_softmax over the bin axis. Used by SONN when param.model.soft_binner=True as an alternative to a learned linear head.

The centers, evenly spaced from 0.05 to 0.95, are a buffer, not a parameter: they move and are saved with the module, but no optimizer sees them, so the passes that switch gradients on for every parameter of the model (Trainer.train_finetune among them) leave them in place.

forward

forward(x: Tensor) -> torch.Tensor

Turn a scalar prediction into logits over the bin centers.

Parameters:

Name	Type	Description	Default
x	(B,) or (B, T) tensor	Prediction(s) on the scale of the centers (0.05 to 0.95).	required

Returns:

Name	Type	Description
logits	(B, n_bins) or (B, n_bins, T) tensor	-scale*(x - center)^2, highest at the nearest center.