calibrax.metrics Registry

Registry infrastructure for metric discovery and management. The singleton MetricRegistry indexes all metrics by name, domain, tier, and mathematical properties (symmetry, invariances, properness).

Types

Metric type definitions: MetricTier enum, MetricProperties, and MetricEntry.

MetricTier

Bases: StrEnum

Computational tier of a metric implementation.

MetricSignature

Bases: StrEnum

Input signature type for a metric function.

Documents what inputs the metric expects, enabling correct dispatch and clear documentation.

MetricProperties(*, is_true_metric: bool = False, is_symmetric: bool = False, is_proper: bool = False, is_differentiable: bool = True, is_jit_compatible: bool = True, invariances: tuple[str, ...] = ()) dataclass

Mathematical and capability properties of a metric.

Groups boolean axiom flags and invariance metadata that describe a metric's mathematical guarantees and JAX compatibility.

Attributes:

Name	Type	Description
is_true_metric	bool	Satisfies identity + symmetry + triangle inequality (metric space axioms).
is_symmetric	bool	d(x, y) = d(y, x). True for all proper metrics, false for KL divergence, etc.
is_proper	bool	Is a proper scoring rule (minimized by true distribution). Relevant for calibration/probabilistic metrics (Brier, log loss, CRPS).
is_differentiable	bool	Can be used inside jax.grad. False for metrics involving argmax, sorting with non-differentiable ops, or string operations.
is_jit_compatible	bool	Can be used inside jax.jit. False for string-based metrics (BLEU, ROUGE) that use Python control flow on variable-length inputs.
invariances	tuple[str, ...]	Transformation groups under which the metric is invariant, following the Erlangen Program. Documents what symmetries the metric preserves. Common values: "translation" (Lp norms), "rotation" (Euclidean), "scale" (cosine, MAPE), "permutation" (Hausdorff, Chamfer), "reparametrization" (Fisher), "affine" (Mahalanobis). Empty tuple means invariances are not documented or the metric has no standard invariance classification.

MetricEntry(*, name: str, fn: Callable[..., float] | None, tier: MetricTier, domain: str = 'general', direction: MetricDirection = MetricDirection.LOWER, description: str = '', required_extra: str = '', signature: MetricSignature = MetricSignature.PREDICTIONS_TARGETS, properties: MetricProperties = MetricProperties()) dataclass

Registry metadata for a metric implementation.

Describes HOW to compute a metric (callable, tier, dependencies). Distinct from MetricDef in core/models.py which describes HOW TO INTERPRET a metric result (direction, unit, priority, group).

Attributes:

Name	Type	Description
name	str	Unique metric identifier (e.g., "mse", "fid", "contrastive_loss").
fn	Callable[..., float] | None	Pure function reference for Tier 0, or None for Tier 1-3.
tier	MetricTier	Computational tier determining the execution pattern.
domain	str	Metric domain (e.g., "general", "image", "text", "audio").
direction	MetricDirection	Whether lower or higher values are better.
description	str	Human-readable description of the metric.
required_extra	str	PyPI extra needed ("" = core, "image" = calibrax[image]).
signature	MetricSignature	Input signature type documenting what arguments the metric expects.
properties	MetricProperties	Mathematical and capability properties of the metric.

Registry

MetricRegistry: singleton registry for metric implementations.

MetricRegistry()

Bases: SingletonRegistry[MetricEntry]

Singleton registry for metric implementations.

Extends SingletonRegistry[MetricEntry] with metric-specific queries (by domain, tier) and auto-registration of built-in metrics.

Usage

registry = MetricRegistry() entry = registry.get("mse") print(entry.tier, entry.domain)

get_function(name: str) -> Callable[..., float]

Retrieve the callable for a Tier 0 metric.

Parameters:

Name	Type	Description	Default
name	str	Metric name.	required

Returns:

Type	Description
Callable[..., float]	The metric's pure function.

Raises:

Type	Description
KeyError	If metric not found.
TypeError	If the metric has no callable (Tier 1-3).

list_by_domain(domain: str) -> list[MetricEntry]

Return all metrics in a given domain.

Parameters:

Name	Type	Description	Default
domain	str	Domain name (e.g., "general", "image").	required

Returns:

Type	Description
list[MetricEntry]	List of matching MetricEntry objects.

list_by_tier(tier: MetricTier) -> list[MetricEntry]

Return all metrics of a given computational tier.

Parameters:

Name	Type	Description	Default
tier	MetricTier	The MetricTier to filter by.	required

Returns:

Type	Description
list[MetricEntry]	List of matching MetricEntry objects.

list_jit_compatible() -> list[MetricEntry]

Return all JIT-compatible metrics.

Returns:

Type	Description
list[MetricEntry]	List of MetricEntry objects where is_jit_compatible is True.

list_true_metrics() -> list[MetricEntry]

Return metrics satisfying metric space axioms.

Returns:

Type	Description
list[MetricEntry]	List of MetricEntry objects where is_true_metric is True.

list_proper_scoring_rules() -> list[MetricEntry]

Return proper scoring rules.

Returns:

Type	Description
list[MetricEntry]	List of MetricEntry objects where is_proper is True.

list_by_invariance(invariance: str) -> list[MetricEntry]

Return metrics invariant under a given transformation group.

Parameters:

Name	Type	Description	Default
invariance	str	Transformation name (e.g., "translation", "rotation").	required

Returns:

Type	Description
list[MetricEntry]	List of MetricEntry objects with the specified invariance.

register_metric(name: str, *, tier: MetricTier = MetricTier.PURE_FUNCTION, domain: str = 'general', direction: MetricDirection = MetricDirection.LOWER, description: str = '', required_extra: str = '', signature: MetricSignature = MetricSignature.PREDICTIONS_TARGETS, properties: MetricProperties | None = None) -> Callable[[Callable[..., float]], Callable[..., float]]

Decorator that registers a function in the MetricRegistry.

Parameters:

Name	Type	Description	Default
name	str	Unique metric name.	required
tier	MetricTier	Computational tier.	PURE_FUNCTION
domain	str	Metric domain.	'general'
direction	MetricDirection	Whether lower or higher is better.	LOWER
description	str	Human-readable description.	''
required_extra	str	PyPI extra needed for this metric.	''
signature	MetricSignature	Input signature type.	PREDICTIONS_TARGETS
properties	MetricProperties | None	Mathematical and capability properties of the metric.	None

Returns:

Type	Description
Callable[[Callable[..., float]], Callable[..., float]]	Decorator that registers the function and returns it unchanged.