metrax.SpearmanRankCorrelation

class metrax.SpearmanRankCorrelation(predictions: Array, labels: Array)

Bases: Metric

Computes the Spearman rank correlation coefficient.

The Spearman rank correlation coefficient measures the monotonic relationship between two variables. It is defined as the Pearson correlation coefficient between the ranked variables.

\[\rho = 1 - \frac{6 \sum d_i^2}{n(n^2 - 1)}\]

where:

• \(d_i\) is the difference between the ranks of each observation

• \(n\) is the number of observations

This implementation accumulates all predictions and labels to compute the exact ranks upon calling compute().

Warning

For very large datasets, this may lead to Out-of-Memory (OOM) errors.

predictions

Accumulated predictions.

Type:

jax.Array

labels

Accumulated labels.

Type:

jax.Array

__init__(predictions: Array, labels: Array) → None

Methods

__init__(predictions, labels)
compute()	Computes final metrics from intermediate values.
compute_value()	Wraps compute() and returns a values.Value.
empty()	Returns an empty instance (i.e. .merge(Metric.empty()) is a no-op).
from_fun(fun)	Calls cls.from_model_output with the return value from fun.
from_model_output(predictions, labels, **kwargs)	Creates a Metric from model outputs.
from_output(name)	Calls cls.from_model_output with model output named name.
merge(other)	Returns Metric that is the accumulation of self and other.
reduce()	Reduces the metric along it first axis by calling _reduce_merge().
replace(**updates)	Returns a new object replacing the specified fields with new values.

Attributes

predictions
labels

predictions: Array

labels: Array

classmethod empty() → SpearmanRankCorrelation

Returns an empty instance (i.e. .merge(Metric.empty()) is a no-op).

classmethod from_model_output(predictions: Array, labels: Array, **kwargs) → SpearmanRankCorrelation

Creates a Metric from model outputs.

merge(other: SpearmanRankCorrelation) → SpearmanRankCorrelation

Returns Metric that is the accumulation of self and other.

Parameters:

other – A Metric whose intermediate values should be accumulated onto the values of self. Note that in a distributed setting, other will typically be the output of a jax.lax parallel operator and thus have a dimension added to the dataclass returned by .from_model_output().

Returns:

A new Metric that accumulates the value from both self and other.

compute() → Array

Computes final metrics from intermediate values.

__init__(predictions: Array, labels: Array) → None

replace(**updates)

Returns a new object replacing the specified fields with new values.