Trainable Masks: Freezing Leaves

Trainability is encoded as a boolean PyTree mask matching the predictors pytree's structure. The training loop calls eqx.partition(predictors, mask) once at start, optimises only the True leaves, and re-combines.

trainable_mask builds the default mask (every inexact-array leaf trainable). The freeze_* helpers compose to zero out subsets — by path, by module type, or by arbitrary predicate.

Quick links

• default_trainable
• trainable_mask
• freeze_paths
• freeze_modules_of_type
• freeze_where
• frozen_default_mask
• count_trainable_params

---

default_trainable()

from jaxhybridmodels.trainable import default_trainable  ·  also re-exported as jaxhybridmodels.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>,
) -> Any

Build a boolean mask matching the structure of predictors.

Applies predicate to every leaf, returning a tree of the same shape whose leaves are bool. Both optimisers take the result unchanged.

Source

---

freeze_paths()

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

freeze_paths(mask: 'Any', paths: 'tuple[str, ...]') -> 'Any'

Return a new mask with leaves at paths set to False.

Path syntax is dot-joined segments addressing the mask PyTree from its root. Each segment is the bare key produced by jax.tree_util.tree_flatten_with_path: attribute names for eqx.Module fields, integer indices for tuples and lists, and string keys for dicts. Example: "inner.mlp.layers.0.weight" addresses mask.inner.mlp.layers[0].weight.

A path matching nothing raises, listing the closest real paths. Ignoring it quietly would leave a leaf the caller believed frozen training as normal, which shows up as a wrong experiment, not a wrong program.

Source

---

freeze_modules_of_type()

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

freeze_modules_of_type(mask: 'Any', predictors: 'Any', cls: 'type') -> 'Any'

Return a new mask with every leaf inside any subtree of type cls set to False.

Walks mask and predictors in lockstep. When a node in predictors is an instance of cls, the matching sub-mask is replaced wholesale by an all-False subtree.

The common use is freeze_modules_of_type(mask, predictors, BoundScaler), which freezes every scaler's temperature. The temperature sets how sharply the squash saturates and is not meant to drift while the model trains.

Source

---

freeze_where()

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

freeze_where(
    mask: 'Any',
    predictors: 'Any',
    fn: 'Callable[[eqx.Module], bool]',
) -> Any

Return a new mask with every leaf inside any submodule satisfying fn set to False.

fn must be a structural test, such as an isinstance check or a look at a static field. It must not compare leaf values. The walk pairs a mask node, whose leaves are booleans, with a predictors node, whose leaves are arrays, so a value comparison has no defined meaning here.

Source

---

frozen_default_mask()

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

frozen_default_mask(predictors: 'Any', *classes: 'type') -> 'Any'

The default mask with every leaf of classes frozen.

Shortcut for the composition every example writes by hand::

mask = trainable_mask(predictors)
mask = freeze_modules_of_type(mask, predictors, BoundScaler)

Freezing BoundScaler leaves (their temperature) is the common case, so frozen_default_mask(predictors, BoundScaler) is the conventional starting mask for a hybrid ODE fit.

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') -> 'int'

Number of trainable scalar parameters selected by mask.

Sums the sizes of every leaf the mask marks True. Useful for reporting the effective search dimension before a run (e.g. to sanity check an evosax budget or a phase-transition threshold).

Source