Utilities Reference¤
This page contains the API reference for all utility functions and classes throughout the library.
Loss Utilities¤
Filter Functions¤
In order to filter fixed loss weights during training updates, there is a utility function to pass to the
equinox.filter_... or equinox.partition functions as a filter_spec:
neojax.losses.utils.is_learnable_loss_weight(model: PyTree) -> PyTree
Creates a filter spec to correctly handle learnable weights.
Produces a boolean PyTree that can be used directly with
eqx.partition, eqx.filter_grad, or eqx.filter_jit.
It returns False for most fields, but True for:
1. All inexact arrays (trainable parameters) in the model.
2. Loss weight attrs only if their learnable_weight is True.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
PyTree
|
The model or list of losses to generate a filter for. |
required |
Returns:
| Type | Description |
|---|---|
PyTree
|
A boolean PyTree of the same structure as |
Example
import equinox as eqx
from neojax.losses.utils import is_learnable_loss_weight
# Directly use as the filter for grad
filter_spec = is_learnable_loss_weight(model)
grads = eqx.filter_grad(loss_fn, filter=filter_spec)(model, ...)