calibrax.core.adapters

Adapter classes for wrapping external objects (ML models, frameworks) to work with Calibrax's interfaces. Includes BenchmarkAdapter (ABC for non-JAX targets) and NNXBenchmarkAdapter (inherits nnx.Module for JIT compatibility).

Generic adapters for the calibrax benchmarking framework.

Provides BenchmarkAdapter ABC for non-NNX targets, NNXBenchmarkAdapter (nnx.Module) for JIT-compatible NNX wrapping, and AdapterRegistry for managing adapter resolution.

These base classes are intentionally thin — they provide only identity, target access, and discoverability. Domain-specific methods (predict, sample, iterate, solve, etc.) are added by sister repos that extend these bases.

BenchmarkAdapter(target)

Bases: ABC

Base class for non-NNX benchmark adapters.

Wraps an arbitrary target (model, data pipeline, solver, etc.) with identity and discoverability. Subclasses add domain-specific methods — this base imposes no interface beyond name, target, and can_adapt.

For NNX models, use NNXBenchmarkAdapter instead — it inherits from nnx.Module for JIT/vmap/grad compatibility.

Initialize the adapter with a target.

Resolves the name from target.name, then target.model_name, falling back to "unknown".

Parameters:

Name	Type	Description	Default
target	Any	The object to adapt for benchmarking.	required

target property

Get the wrapped target object.

Returns:

Type	Description
Any	The original target passed to the constructor.

name property

Get the name of the adapted target.

Returns:

Type	Description
str	The target name string.

can_adapt(target) classmethod

Check if this adapter can handle the given target.

Returns False by default — subclasses override with specific checks.

Parameters:

Name	Type	Description	Default
target	object	The object to check.	required

Returns:

Type	Description
bool	True if this adapter can wrap the target.

NNXBenchmarkAdapter(model)

Bases: Module

JIT-compatible adapter for Flax NNX modules.

Inherits from nnx.Module so it participates in NNX's graph system. This enables:

• nnx.jit for JIT-compiled execution
• nnx.vmap for batched execution
• Proper state mutation tracking (RNG, batch norm, etc.)
• nnx.split / nnx.merge for state serialization

Like BenchmarkAdapter, this base is intentionally thin — subclasses add domain-specific methods (predict, sample, solve, etc.).

Initialize the adapter with an NNX module.

The model becomes a tracked sub-module in the NNX graph system, ensuring its parameters and state are handled correctly during JIT compilation and other transforms.

Resolves the name from model.name, then model.model_name, falling back to "unknown".

Parameters:

Name	Type	Description	Default
model	Module	The NNX module to adapt.	required

name property

Get the name of the adapted model.

Returns:

Type	Description
str	The model name string.

can_adapt(target) classmethod

Check if the target is a Flax NNX Module.

Parameters:

Name	Type	Description	Default
target	object	The object to check.	required

Returns:

Type	Description
bool	True if the target is an nnx.Module instance.

AdapterRegistry()

Registry of benchmark adapters.

Maintains an ordered list of adapter classes and resolves the appropriate adapter for a given target via can_adapt checks.

Accepts both BenchmarkAdapter subclasses (for non-NNX targets) and NNXBenchmarkAdapter subclasses (nnx.Module-based, for JIT-compatible NNX adapters).

Initialize an empty adapter registry.

register(adapter_cls)

Register an adapter class (highest priority first).

The adapter class must have a can_adapt(target) classmethod and accept a target as its first constructor argument.

Parameters:

Name	Type	Description	Default
adapter_cls	type	The adapter class to register.	required

adapt(target)

Find and apply a suitable adapter for the target.

Parameters:

Name	Type	Description	Default
target	Any	The object to adapt.	required

Returns:

Type	Description
Any	An adapter wrapping the target.

Raises:

Type	Description
ValueError	If no registered adapter can handle the target.

reset()

Remove all registered adapters.

adapt(target)

Adapt a target using the default registry.

Parameters:

Name	Type	Description	Default
target	Any	The object to adapt.	required

Returns:

Type	Description
Any	An adapter wrapping the target.

Raises:

Type	Description
ValueError	If no adapter can handle the target.

register_adapter(adapter_cls)

Register an adapter class into the default registry.

Parameters:

Name	Type	Description	Default
adapter_cls	type	The adapter class to register.	required