pyproject.toml

[project]
name = "jax_loop_utils"
version = "0.0.12"
description = "Libraries for ML training loops in JAX."
readme = "README.md"
requires-python = ">=3.11"
license = { text = "Apache 2.0" }
authors = [{ name = "Astera Institute", email = "no-reply@astera.org" }]
dependencies = [
    "absl-py",
    "etils[epath,epy]",
    "jax",
    "numpy",
    "packaging",
    "Pillow",
    "wrapt",
]
classifiers = [
    "Development Status :: 4 - Beta",
    "Intended Audience :: Developers",
    "Intended Audience :: Science/Research",
    "License :: OSI Approved :: Apache Software License",
    "Programming Language :: Python",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
]
keywords = ["JAX", "machine learning"]

[project.urls]
Homepage = "http://github.com/Astera-org/jax_loop_utils"

[project.optional-dependencies]
mlflow = ["mlflow-skinny>=2.0"]
pyright = ["pyright"]
# for synopsis.ipynb
synopsis = ["chex", "flax", "ipykernel", "matplotlib"]
tensorflow = ["tensorflow>=2.12"]
test = ["chex", "pytest", "pytest-cov"]
torch = ["torch>=2.0"]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build]
exclude = ["*_test.py"]

[tool.uv]

# restrict to platforms we care about so that version resolution is faster and more likely to succeed
# (e.g. don't fail if a package isn't built for windows)
environments = [
    "sys_platform == 'darwin' and platform_machine == 'arm64'",
    "sys_platform == 'linux' and platform_machine == 'aarch64'",
    "sys_platform == 'linux' and platform_machine == 'x86_64'",
]

[tool.pytest.ini_options]
filterwarnings = [
    # Please only ignore warnings that come from a transitive dependency that we
    # can't easily avoid.
    # See https://docs.pytest.org/en/stable/how-to/capture-warnings.html#controlling-warnings
    # action:message:category:module:line
    "error",
]