Parameter containers

This page lists exact signatures. Four related documents cover different ground:

Document	What it gives you
Parameters, transforms, and regularization	The guided tour, with three worked models
Choose a parameter transform	The transform catalog, by constrained domain
Constrain and regularize parameters	Short task recipes
The parameter model	The why behind the design

1. Parameter containers

Param wraps a trainable array with an optional bijective transform for constrained optimization and an optional regularization penalty. The value passed to the constructor is interpreted in constrained space; the inverse transform derives the unconstrained array that is actually stored and updated.

Const is the non-trainable counterpart: it participates in the forward pass but is excluded from ParamState collection, so optimizers and grad never see it.

Param	A module has neural network parameters for optional transform and regularization.
Const	A module has non-trainable constant parameter.

2. Model-level parameter API

These methods live on Module and operate across the whole module graph: discovering parameters, aggregating their penalties, and scoping their caches.

Important

All four methods accept only allowed_hierarchy=(min_depth, max_depth) — a closed interval, where a Param attached directly to the root counts as depth 1. They do not accept a filter function. Filter the result with an ordinary list comprehension:

regularized = [(n, p) for n, p in model.named_param_modules() if p.reg is not None]
trainable   = [p for p in model.param_modules() if p.fit]

Traversal includes Const, since it is a Param subclass; check .fit when trainability is what you mean.

Module.param_modules	Collect all Param parameters in this module and children.
Module.named_param_modules	Iterate over (name, parameter) pairs.
Module.reg_loss	Compute total regularization loss from all Param parameters.
Module.param_precompute	Context manager to temporarily cache all Param parameters.

The per-parameter counterparts on Param itself:

Param.value	Get current parameter value after applying transform.
Param.set_value	Set parameter value from constrained space.
Param.reg_loss	Calculate regularization loss.
Param.cache	Manually cache the transformed value.
Param.clear_cache	Explicitly clear the parameter transformation cache.
Param.cache_stats	Get cache statistics (for debugging/monitoring).
Param.clip	Clamp parameter value in-place.
Param.reset_to_prior	Reset parameter value to regularization prior value.
Param.init	Initialize parameters.

Note

Param.value() does not populate the cache. The cache is opt-in: only Param.cache() or Module.param_precompute() warms it. This prevents a value traced inside jit from being cached and reused outside that trace.

3. Parameter transforms

Bijections between an unconstrained space and a constrained one, letting an optimizer work without walls while the constraint holds by construction. All implement forward(), inverse(), and optionally log_abs_det_jacobian() for probabilistic applications. Pass one to Param as t=.

Base class and the no-op default:

Transform	Abstract base class for bijective parameter transformations.
IdentityT	Identity transformation (no-op).

Positive and negative half-lines:

SoftplusT	Softplus transformation mapping unbounded values to positive semi-infinite interval.
ExpT	Exponential transformation mapping (-inf, +inf) to (lower, +inf).
PositiveT	Transformation constraining parameters to be strictly positive (0, +∞).
ReluT	ReLU transform with lower bound: forward(x) = relu(x) + lower_bound
NegSoftplusT	Negative softplus transformation mapping unbounded values to negative semi-infinite interval.
NegativeT	Transformation constraining parameters to be strictly negative (-∞, 0).
LogT	Log transformation mapping (lower, +inf) to (-inf, +inf).

Bounded intervals:

SigmoidT	Sigmoid transformation mapping unbounded values to a bounded interval.
ScaledSigmoidT	Sigmoid transformation with adjustable sharpness/temperature.
ClipT	Transformation with clipping to specified bounds.
TanhT	Tanh-based transformation mapping (-inf, +inf) to (lower, upper).
SoftsignT	Softsign-based transformation mapping (-inf, +inf) to (lower, upper).

Structured domains:

SimplexT	Stick-breaking transformation for simplex constraint.
UnitVectorT	Transformation to unit vectors (L2 norm = 1).
OrderedT	Transformation ensuring ordered (monotonically increasing) output.

Reparameterizations and composition:

AffineT	Affine (linear) transformation with scaling and shifting.
PowerT	Power (Box-Cox) transformation for stabilizing variance.
ChainT	Composition of multiple transformations applied sequentially.
MaskedT	Selective transformation using a boolean mask.

4. Standard regularizations

Classical penalties encouraging sparsity, smoothness, or structural properties such as orthogonality and bounded spectral norm. Pass one to Param as reg=; Module.reg_loss() sums them across the model.

Regularization	Abstract base class for parameter regularization.
L1Reg	L1 (Lasso) regularization.
L2Reg	L2 (Ridge) regularization.
ElasticNetReg	Elastic Net regularization (combination of L1 and L2).
HuberReg	Huber regularization (robust regularization).
GroupLassoReg	Group Lasso regularization.
TotalVariationReg	Total Variation regularization.
MaxNormReg	Max Norm regularization (soft constraint).
EntropyReg	Entropy regularization.
OrthogonalReg	Orthogonal regularization.
SpectralNormReg	Spectral Norm regularization.
ChainedReg	Composite regularization that chains multiple regularizations together.

5. Prior distribution-based regularizations

Probabilistic penalties derived from a prior distribution, for Bayesian-inspired parameter estimation: variational inference, maximum a posteriori estimation, and uncertainty quantification. Each contributes the negative log-density of the parameter under its prior, and implements loss(), sample_init(), and reset_value() for prior-based initialization.

GaussianReg	Gaussian prior regularization.
StudentTReg	Student's t-distribution prior regularization.
CauchyReg	Cauchy prior regularization.
UniformReg	Uniform prior regularization (soft bounded constraint).
BetaReg	Beta prior regularization (for parameters in [0, 1]).
LogNormalReg	Log-normal prior regularization (for positive parameters).
ExponentialReg	Exponential prior regularization (for positive parameters).
GammaReg	Gamma prior regularization (for positive parameters).
InverseGammaReg	Inverse-Gamma prior regularization (for variance parameters).
LogUniformReg	Log-uniform (Jeffreys) prior regularization (scale-invariant).
HorseshoeReg	Horseshoe prior regularization (strong sparsity with heavy tails).
SpikeAndSlabReg	Spike-and-slab prior regularization (variable selection).
DirichletReg	Dirichlet prior regularization (for probability simplexes).

6. Use cases

Common patterns, each a minimal snippet you can lift into a model. All follow the same rule: .value() returns the parameter in its constrained domain (use it in the forward pass), while .val is the underlying ParamState the optimizer updates in unconstrained space.

A positive-only rate or time constant. SoftplusT keeps the value above a floor no matter where the optimizer lands, so a membrane time constant or a firing rate never goes negative:

tau = Param(jnp.array(20.0), t=SoftplusT(lower=1.0))   # ms, stays > 1.0
dv = (-v + i_input) / tau.value()

A mixing weight or gate bounded to [0, 1]. SigmoidT maps the whole real line into an open interval, so a convex-combination coefficient cannot leave its range:

alpha = Param(jnp.array(0.5), t=SigmoidT(lower=0.0, upper=1.0))
blended = alpha.value() * fast + (1.0 - alpha.value()) * slow

A learned categorical distribution. SimplexT guarantees .value() is a valid probability vector (non-negative, sums to one) for any unconstrained input — handy for a learned prior over components or a soft attention/routing weight:

weights = Param(jnp.zeros(4), t=SimplexT())
mixture = jnp.tensordot(weights.value(), components, axes=1)

Weight decay and sparsity. Attach a penalty with reg= and add .reg_loss() to the data loss. L2Reg shrinks weights smoothly; L1Reg drives them to exact zeros:

dense = Param(random.randn(din, dout) * 0.1, reg=L2Reg(weight=1e-4))
gate  = Param(random.randn(dout), reg=L1Reg(weight=1e-3))

A Bayesian / MAP prior on a parameter. A prior-based regularization contributes the negative log-density of the parameter under a prior, turning training into maximum-a-posteriori estimation:

mu = Param(jnp.zeros(dout), reg=GaussianReg(mean=0.0, std=1.0))

Aggregating penalties across a whole model. Module.reg_loss() sums every penalty in the module tree in one call, so the training step never has to enumerate parameters by hand:

def loss_fn(batch):
    data_loss = mse(model(batch.x), batch.y)
    return data_loss + model.reg_loss()

Warming caches for constrained parameters. When many parameters share expensive transforms, call Module.param_precompute() once outside the hot loop; .value() then reads the cached constrained array instead of recomputing the bijection on every forward pass:

model.param_precompute()      # warm every Param cache in the tree
for batch in loader:
    predictions = model(batch.x)