Trainable Masks: Freezing Leaves

Trainability is a boolean pytree mask with the structure of the predictors. Training partitions the predictors on it once and optimises only the True leaves.

trainable_mask builds the default mask (every inexact-array leaf except Constants and scalers). freeze_mask is what fit(freeze=...) calls: it sets the leaves matched by paths, classes or predicates to False, on top of an optional mask=. Pass its result as fit(trainable=...) when freeze= is not enough.

Quick links

• default_trainable
• trainable_mask
• freeze_mask
• count_trainable_params

---

default_trainable()

from jaxhybridmodels.trainable import default_trainable

default_trainable(leaf: 'Any') -> 'bool'

Default trainability rule. True only for inexact-array leaves.

"Inexact" means a JAX array with a float or complex dtype. Everything else is fixed, including ints, bools, Python scalars and static-field values, which is what a gradient-based optimiser can actually update.

Source

---

trainable_mask()

from jaxhybridmodels.trainable import trainable_mask  ·  also re-exported as jaxhybridmodels.trainable_mask

trainable_mask(
    predictors: 'Any',
    predicate: 'Callable[[Any], bool]' = <default_trainable>,
    train_scalers: 'bool' = False,
) -> Any

Build a boolean mask matching the structure of predictors.

Applies predicate to parameter leaves, returning a tree of the same shape whose leaves are bool. Constant subtrees are always fixed. BoundScaler subtrees are fixed unless train_scalers is set: a trainable temperature moves gradient between the scaler and the inner network and tends to slow convergence. Both optimisers take the result unchanged.

Source

---

freeze_mask()

from jaxhybridmodels.trainable import freeze_mask

freeze_mask(
    predictors: 'Any',
    freeze: 'Sequence[FreezeSpec] | FreezeSpec',
    mask: 'Any' = None,
) -> Any

The mask behind fit(freeze=...): mask with the freeze entries set to False.

mask defaults to trainable_mask of predictors. Each freeze entry is one of

• a dotted path naming a leaf or a whole subtree: "rate" for predictors["rate"], "rate.inner.mlp.layers.0.weight" for one leaf. Segments are the bare keys of jax.tree_util.tree_flatten_with_path: attribute names for eqx.Module fields, integer indices for tuples and lists, and dict keys;
• a module class, freezing every instance (hm.Constant is already frozen; hm.MLPPredictor is the usual one);
• a predicate (module) -> bool, called on the real modules of predictors, so it may read their values.

A single entry may be passed bare. An entry that matches no leaf raises, listing the closest real paths: ignoring it would leave a leaf the caller believed frozen training as normal, which shows up as a wrong experiment rather than a wrong program.

Examples

::

mask = freeze_mask(predictors, ("offset", hm.MLPPredictor))
# Chain onto a mask that also trains the scalers.
base = trainable_mask(predictors, train_scalers=True)
mask = freeze_mask(predictors, "offset", mask=base)

Source

---

count_trainable_params()

from jaxhybridmodels.trainable import count_trainable_params  ·  also re-exported as jaxhybridmodels.count_trainable_params

count_trainable_params(predictors: 'Any', *, mask: 'Any' = None) -> 'int'

Number of trainable scalar parameters mask selects in predictors.

mask defaults to trainable_mask of predictors. Sums the sizes of every leaf the mask marks True. Useful for reporting the effective search dimension before a run, say to check an Evosax budget.

Source