PyPI page
Author:
Mahdi Shafiei
License:
MIT License
Copyright (c) 2026 Mahdi Shafiei
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated docu...
Summary:
The training plumbing a JAX/Flax NNX project otherwise writes for itself: process setup, device sharding, checkpoints, early stopping, logging and remote GPU runs
Latest version:
0.1.22
Required dependencies:
etils
|
flax
|
jax
|
jaxlib
|
numpy
|
optax
|
orbax-checkpoint
|
pydantic
Optional dependencies:
bandit
|
build
|
coverage
|
deptry
|
diff-cover
|
griffe
|
hypothesis
|
import-linter
|
interrogate
|
jax
|
jax-metal
|
matplotlib
|
mkdocs
|
mkdocs-material
|
mkdocstrings
|
mkdocstrings-python
|
mlflow
|
modal
|
pip-audit
|
pre-commit
|
pydoclint
|
pygments
|
pylint
|
pymdown-extensions
|
pyright
|
pytest
|
pytest-cov
|
pytest-env
|
pytest-json-report
|
pytest-randomly
|
pytest-timeout
|
pytest-xdist
|
pyyaml
|
radon
|
ruff
|
toml-sort
|
twine
|
validate-pyproject
|
vulture
|
wandb
|
xenon
Downloads last day:
173
Downloads last week:
1,738
Downloads last month:
15,855
7-day smoothing