Base Metric

Base Metric

neojax.metrics.base_metric.BaseMetric

Abstract base metric.

New metrics and loss functions should inherit from BaseMetric. This class cannot be instantiated directly. Subclasses must be final following the "Abstract or Final" class design pattern.

Every metric evaluates predictions against a ground-truth target. It can also act as a differentiable loss function during training. To support learnable weighting (e.g., in multi-task loss compositions), each metric carries a weight field that can optionally be flagged as learnable.

Parameters:

Name	Type	Description	Default
weight	float | int	(Learnable) scalar weight coefficient of the metric. The metric is computed as 'metric * weight'. Default is 1.0. If learnable_weight is True, weight has to be strictly positive and will be optimized via gradient descent.	required
learnable_weight	bool	Optional flag indicating whether weight is learnable. If True, this flag is used in the is_learnable_metric_weight filter function to indicate to eqx.filter_... or eqx.partition that weight should be adapted by the optimizer. Default is False.	False

Raises:

Type	Description
ValueError	If learnable_weight is True and weight is negative or zero.

Attributes

• raw_weight (Float[Array, ""]): Scalar weight coefficient for the metric. Used to balance multiple metrics/losses in compositions.
• learnable_weight (bool): Static flag indicating whether the weight is learnable via gradient descent.

Why the raw weight is not learnable

A metric has a non-negative value. If the raw weight was learnable this would give the optimizer an easy way to drive the loss down: reduce the weight to zero or drive the weight negative which rewards the model to fit worse. To prevent such 'cheating', the weight is transformed by jax.nn.softplus which makes it strictly positive. An arbitrary fixed weight can still be set as long as learnable_weight is False.

weight property

The scalar weight. Strictly positive if learnable_weight is True.

__call__(model: equinox._module._module.Module | None = None, *, target: Float[Array, 'b c ...'], x: PyTree[Float[Array, 'b ...']] | None = None, pred: Float[Array, 'b c ...'] | None = None, **kwargs: typing.Any) -> Float[Array, '']

Computes metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The neural operator model being evaluated. Default is None. Required for operator-dependent metrics (e.g., Sobolev metrics).	None
target	Float[Array, 'b c ...']	Ground truth array of shape (batch, channels, *spatial).	required
x	PyTree[Float[Array, 'b ...']] | None	Input array to the model of shape (batch, in_channels, *spatial). Required for operator-dependent metrics.	None
pred	Float[Array, 'b c ...'] | None	Optional pre-computed model prediction array of shape (batch, channels, *spatial). If not provided, the metric will attempt to compute it by evaluating the model on x.	None
**kwargs	Any	Additional parameters for custom metrics.	required

Returns:

Type	Description
Float[Array, '']	Scalar metric value.

---

Custom Metrics

When implementing custom metrics, inherit from BaseMetric.

Each custom metric should implement the __call__ method, taking the predicted array and the target array as inputs and returning a scalar metric.