Metadata-Version: 2.4
Name: xtrax
Version: 0.3.1
Summary: High-performance composable JAX library
Project-URL: Homepage, https://github.com/maraxen/xtrax
Project-URL: Documentation, https://xtrax.readthedocs.io
Project-URL: Repository, https://github.com/maraxen/xtrax
Project-URL: Issues, https://github.com/maraxen/xtrax/issues
Project-URL: Changelog, https://github.com/maraxen/xtrax/blob/main/CHANGELOG.md
Author-email: Marielle Russo <67157875+maraxen@users.noreply.github.com>
License-Expression: Apache-2.0
License-File: LICENSE
Keywords: distributed,equinox,jax,sparse,training
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3.13
Classifier: Typing :: Typed
Requires-Python: >=3.13
Requires-Dist: equinox>=0.11.0
Requires-Dist: grain>=0.2.0
Requires-Dist: jax>=0.4.36
Requires-Dist: jaxlib>=0.4.36
Requires-Dist: numpy>=1.26
Requires-Dist: optax>=0.2.3
Requires-Dist: orbax-checkpoint>=0.6.0
Requires-Dist: pytest-asyncio>=0.23
Provides-Extra: cli
Requires-Dist: flatbuffers>=2; extra == 'cli'
Requires-Dist: tyro<2,>=0.9; extra == 'cli'
Provides-Extra: dev
Requires-Dist: beartype>=0.18.0; extra == 'dev'
Requires-Dist: chex>=0.1.86; extra == 'dev'
Requires-Dist: complexipy>=5.6.1; extra == 'dev'
Requires-Dist: import-linter>=2.1; extra == 'dev'
Requires-Dist: interrogate>=1.7.0; extra == 'dev'
Requires-Dist: jaxlint>=0.1.0a1; extra == 'dev'
Requires-Dist: jaxtyping>=0.2.28; extra == 'dev'
Requires-Dist: libcst>=1.5.0; extra == 'dev'
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
Requires-Dist: pytest-benchmark>=4.0; extra == 'dev'
Requires-Dist: pytest-cov>=5.0; extra == 'dev'
Requires-Dist: pytest>=8.0; extra == 'dev'
Requires-Dist: ruff>=0.4.0; extra == 'dev'
Requires-Dist: ty; extra == 'dev'
Provides-Extra: eda
Requires-Dist: matplotlib>=3.8; extra == 'eda'
Requires-Dist: pandas>=2.0; extra == 'eda'
Requires-Dist: seaborn>=0.13; extra == 'eda'
Description-Content-Type: text/markdown

# xtrax

[![PyPI - Version](https://img.shields.io/pypi/v/xtrax.svg)](https://pypi.org/project/xtrax/)
[![Tests](https://github.com/maraxen/xtrax/actions/workflows/ci.yml/badge.svg)](https://github.com/maraxen/xtrax/actions)
[![Docs](https://img.shields.io/readthedocs/xtrax.svg)](https://xtrax.readthedocs.io)
[![Coverage](https://img.shields.io/badge/coverage-tier1%2095%25-brightgreen)](https://github.com/maraxen/xtrax)
[![License: Apache 2.0](https://img.shields.io/badge/license-Apache%202.0-blue.svg)](https://github.com/maraxen/xtrax/blob/main/LICENSE)

A set of composable building blocks for JAX/Equinox training loops — engine + trainer orchestration, safety-checked steps, axis tiling strategies, inference-time sparsification, distributed/sharding helpers, streaming output callbacks, and orbax checkpointing — extracted from the author's research code.

## Status

**xtrax is alpha, experimental software** built primarily for the author's personal research use. APIs may change without notice between releases; no backward-compatibility guarantees pre-1.0. Issues and pull requests are welcome, but support is best-effort — the project exists first to serve the author's own JAX training workflows.

## Why xtrax?

- **Composable training steps** — Use `Trainer` or `SafetyTrainStep` with your own loss functions and optimizers
- **Safety-checked arithmetic** — Avoid NaN/Inf propagation with safe operations (`safe_norm`, `safe_reciprocal`)
- **Flexible tiling strategies** — Partition computations with `AxisSpec`, `BatchPlan`, and `Vmap`/`SafeMap` for data and model parallelism
- **Inference sparsification** — Apply structured sparsity masks with `SparseConfig` and `sparsify_model`
- **Distributed helpers** — Initialize and coordinate multi-GPU/TPU training with `init_dist`, `LogicalMesh`, and sharding utilities

## Installation

```bash
pip install xtrax
```

Requires Python 3.13 or later.

## Quick Start

```python
import jax
import optax
from xtrax import Trainer, ResumableState, Engine, save_checkpoint, load_checkpoint

# 1. Create a simple loss function
def loss_fn(model, batch):
    predictions = model(batch["inputs"])
    return jnp.mean((predictions - batch["targets"]) ** 2)

# 2. Set up trainer with optimizer
optimizer = optax.adam(1e-3)
trainer = Trainer(loss_fn=loss_fn, optimizer=optimizer)

# 3. Initialize training state
model = ...  # Your equinox model
opt_state = optimizer.init(...)
state = ResumableState(model=model, opt_state=opt_state, step=0)

# 4. Run a training step
new_state, metrics = trainer.step(state, batch={"inputs": x, "targets": y})
print(f"Loss: {metrics['loss']}")
```

For a complete training loop with callbacks and checkpointing, use the `Engine`:

```python
from xtrax import Engine, DataModule

# Create or load a DataModule (must implement train_iter())
data = DataModule(...)

# Create an engine with trainer and optional callbacks
engine = Engine(trainer=trainer)

# Run multi-epoch training with checkpoint saving
final_state = engine.fit_sync(
    state=state,
    data=data,
    num_epochs=10,
    checkpoint_dir="./checkpoints"
)
```

## Getting Results Out

### Streaming Callbacks

Log metrics asynchronously to files or external services:

```python
from xtrax.io import BoundedCallbackHandler, async_indexed_stream

# Create a custom async callback
class LogCallback:
    async def on_step_end(self, state, metrics):
        print(f"Step {state.step}: {metrics}")

# Use in your Engine
engine = Engine(
    trainer=trainer,
    callbacks=[LogCallback()]
)
```

### Checkpoint Save and Load

Save model state and restore for inference or resumption:

```python
from xtrax import save_checkpoint, load_checkpoint

# After training
save_checkpoint(checkpoint_dir="./checkpoints/final", state=final_state)

# Load for inference
restored_state = load_checkpoint(checkpoint_dir="./checkpoints/final")
model = restored_state.model

# Run inference
predictions = model(test_inputs)
```

## Documentation

Full API docs, architecture guides, and advanced examples at [xtrax.readthedocs.io](https://xtrax.readthedocs.io).

## Project links

- [GitHub repository](https://github.com/maraxen/xtrax)
- [Issue tracker](https://github.com/maraxen/xtrax/issues)
- [Changelog](CHANGELOG.md)
- [Contributing guide](CONTRIBUTING.md)
- [Citation metadata](CITATION.cff)

## License

Licensed under the [Apache License 2.0](https://github.com/maraxen/xtrax/blob/main/LICENSE).

## Citation

If you use xtrax in research, please cite it:

```bibtex
@software{xtrax,
  title = {xtrax: High-Performance Composable JAX Training},
  author = {Russo, Marielle},
  version = {0.3.0},
  year = {2026},
  url = {https://github.com/maraxen/xtrax}
}
```
