Regression Metrics

Regression Metrics

neojax.metrics.regression_metrics.MSEMetric

Mean Squared Error (MSE) metric.

Computes the mean squared difference between predictions and targets:

\[ \text{MSE}(y, \hat{y}) = \frac{1}{N} \sum_{i=1}^N (y_i - \hat{y}_i)^2 \]

where \(y\) is the ground truth target and \(\hat{y}\) is the model prediction.

Parameters:

Name	Type	Description	Default
weight	float	(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.	1.0
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

Internal Attributes

These fields store the internal state of the metric.

• weight (Float[Array, ""]): Learnable metric weight. Filter during training to prevent updates.
• learnable_weight (bool): Flag indicating whether weight is learnable.

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

Computes the Mean Squared Error (MSE) metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...']	Ground truth array shaped (batch, c, d1, ..., dN).	required
x	Float[Array, 'b in_c ...'] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array. If None, the metric computes vmap(model)(x).	None

Returns:

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

Raises:

Type	Description
ValueError	If pred is None and either model or x is None.

---

neojax.metrics.regression_metrics.RMSEMetric

Root Mean Squared Error (RMSE) metric.

Computes the square root of the mean squared difference between predictions and targets:

\[ \text{RMSE}(y, \hat{y}) = \sqrt{\frac{1}{N} \sum_{i=1}^N (y_i - \hat{y}_i)^2 + \epsilon} \]

where \(y\) is the ground truth target, \(\hat{y}\) is the model prediction, and \(\epsilon = 10^{-7}\) is a small regularization constant to prevent NaN gradients at zero.

Parameters:

Name	Type	Description	Default
weight	float	(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.	1.0
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

Internal Attributes

These fields store the internal state of the metric.

• weight (Float[Array, ""]): Learnable metric weight. Filter during training to prevent updates.
• learnable_weight (bool): Flag indicating whether weight is learnable.

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

Computes the Root Mean Squared Error (RMSE) metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...']	Ground truth array shaped (batch, c, d1, ..., dN).	required
x	Float[Array, 'b in_c ...'] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array. If None, the metric computes vmap(model)(x).	None

Returns:

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

Raises:

Type	Description
ValueError	If pred is None and either model or x is None.

---

neojax.metrics.regression_metrics.R2Metric

Coefficient of Determination (\(R^2\) score) metric.

Computes the proportion of the variance in the target that is predictable from the model prediction:

\[ R^2(y, \hat{y}) = 1 - \frac{\sum_{i=1}^N (y_i - \hat{y}_i)^2}{\sum_{i=1}^N (y_i - \bar{y})^2 + \epsilon} \]

where \(y\) is the ground truth target, \(\hat{y}\) is the model prediction, \(\bar{y}\) is the mean of the ground truth target, and \(\epsilon = 10^{-7}\) is a small regularization constant to prevent division by zero.

Parameters:

Name	Type	Description	Default
weight	float	(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.	1.0
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

Internal Attributes

These fields store the internal state of the metric.

• weight (Float[Array, ""]): Learnable metric weight. Filter during training to prevent updates.
• learnable_weight (bool): Flag indicating whether weight is learnable.

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

Computes the Coefficient of Determination (\(R^2\) score) metric.

Parameters:

Name	Type	Description	Default
model	Module | None	The model being trained. Optional if pred is given.	None
target	Float[Array, 'b c ...']	Ground truth array shaped (batch, c, d1, ..., dN).	required
x	Float[Array, 'b in_c ...'] | None	Model input array shaped (batch, in_c, d1, ..., dN).	None
pred	Float[Array, 'b c ...'] | None	Model prediction array. If None, the metric computes vmap(model)(x).	None

Returns:

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

Raises:

Type	Description
ValueError	If pred is None and either model or x is None.