Install & Compatibility
Where this runs
tested against v0.0.12 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
muslpy 3.10–3.95 runs
build_error
glibcpy 3.10–3.95 runs
installs and imports cleanly · install 17.9s · import 0.000s · 682MB
711MB installed
● package 711MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
metric_writers
✓ import clu.metric_writers
✗ import clu.metric_writers
This quickstart demonstrates how to define and use `clu.metrics` to track custom metrics like accuracy and average loss over multiple training steps in a JAX environment. It shows how to create a `metrics.Collection`, gather outputs from a simulated `train_step`, merge them, and then compute the final aggregate results.
import jax
import jax.numpy as jnp
from clu import metrics
# Define a simple custom metric
class Accuracy(metrics.Metric):
num_correct: metrics.Sum.from_output('correct')
num_total: metrics.Sum.from_output('total')
def compute(self):
return self.num_correct / self.num_total
# Simulate a training step
def train_step(params, batch):
# In a real scenario, this would involve model inference and loss calculation
predictions = jnp.array([0.8, 0.1, 0.9]) # Example predictions
labels = jnp.array([1, 0, 1]) # Example true labels
correct = (predictions > 0.5) == labels
return {
'correct': correct.sum(),
'total': correct.size
}
# Initialize metrics
all_metrics = metrics.Collection.create(
accuracy=Accuracy,
loss=metrics.Average.from_output('batch_loss')
)
# Simulate first batch
params_dummy = {'w': jnp.array([1.0])}
batch_dummy = {}
step_outputs = train_step(params_dummy, batch_dummy)
step_outputs['batch_loss'] = jnp.array(0.1) # Example loss
all_metrics = all_metrics.merge(all_metrics.empty()).gather_from_model_output(**step_outputs)
# Simulate another batch
step_outputs_2 = train_step(params_dummy, batch_dummy)
step_outputs_2['batch_loss'] = jnp.array(0.05) # Example loss
all_metrics = all_metrics.merge(all_metrics.gather_from_model_output(**step_outputs_2))
# Compute and print results
results = all_metrics.compute()
print(f"Average Accuracy: {results['accuracy']:.2f}")
print(f"Average Loss: {results['loss']:.2f}")
Upgrade
Version history
0.0.12latest on PyPI · released Apr 10, 2024
Audit
Dependencies
jaxrequiredCore dependency for JAX-based ML workflows.
tensorflow-datasetsoptionalCommonly used for data loading and input pipelines, especially with `clu.deterministic_data`.