clu

What it is and what it does

CLU is a set of libraries built on top of JAX for structuring machine learning training loops. It abstracts common tasks—metrics aggregation, checkpointing, and loop orchestration—into reusable components so that training code stays readable and concise without sacrificing the flexibility needed for research. The package depends on JAX, Flax, ml-collections, and several utility libraries (absl-py, etils, numpy, wrapt, packaging, typing-extensions) to provide a cohesive toolkit.

The library is designed for researchers and practitioners writing custom training loops in JAX who want to avoid boilerplate while keeping full control over loop logic. It is actively maintained (last commit 2026-07-07) and widely adopted in the JAX and machine learning community. The project is not currently accepting external contributions but can be forked for custom extensions.

Use it for:

• Collecting and aggregating metrics (loss, accuracy, etc.) during JAX training without writing custom metric logic.
• Managing checkpoints and model state during long-running training jobs in JAX.
• Structuring a custom training loop in Flax or raw JAX while reusing common patterns from CLU.
• Prototyping ML experiments in JAX with minimal boilerplate for loop orchestration and logging.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

CLU provides utilities and abstractions for writing machine learning training loops in JAX, including metrics collection, checkpointing, and common patterns to keep loop code concise while retaining research flexibility.

Yes, if you are writing training loops in JAX or Flax. CLU is actively maintained, has no security vulnerabilities, low install friction, and is widely used in the JAX ecosystem. The permissive Apache 2.0 license poses no restrictions. Verify that CLU's feature set matches your specific checkpoint and metrics needs before committing.

Install

clu on PyPI

pip install clu

uv add clu

poetry add clu

Installing clu

Before you install

Low friction: pure Python wheel with no compiled dependencies. Active maintenance—last commit 2026-07-07, though the latest release was 2024-04-10, suggesting the package is stable but not under rapid iteration.

License in practice

Apache 2.0 is permissive; you can use, modify, and distribute CLU freely in commercial and private projects, provided you retain the license notice.

Quickstart

pip install clu

import clu
from clu import metrics

# Use CLU metrics and utilities in your JAX training loop
metric_collection = metrics.Collection()

Requires JAX and jaxlib installed; CLU is a library for JAX-based training, not a standalone tool.

Verify before relying

• Specific metrics types and checkpoint formats supported by CLU beyond what the excerpt describes.
• Whether CLU is actively accepting bug fixes or is in maintenance-only mode given the time since last release.

Package facts

Evidence: clu-0.0.12-py3-none-any.whl

Keywords: JAX, machine, learning

Tags

More Artificial Intelligence packages