pyproject.toml

[project]
name = "jax-ai-stack"
version = "2025.1.9"  # keep in sync with jax_ai_stack/__init__.py
description = ""
readme = "README.md"
license = {file = "LICENSE"}
authors = [{name = "jax-ai-stack authors", email="jax-dev+jax-ai-stack@google.com"}]
classifiers = [
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3 :: Only",
    "License :: OSI Approved :: Apache Software License",
    "Intended Audience :: Science/Research",
]
keywords = []

# pip dependencies of the project
dependencies = [
    "chex==0.1.88",
    "jax==0.4.38",
    "flax==0.10.2",
    "ml_dtypes==0.4.0",
    "optax==0.2.4",
    "orbax-checkpoint==0.11.0",
    "orbax-export==0.0.6",
]

[project.urls]
homepage = "https://github.com/jax-ml/jax_-ai-stack"
repository = "https://github.com/jax-ml/jax-ai-stack"
# Other: `documentation`, `changelog`

[project.optional-dependencies]
# Development deps (unittest, linting, formating,...)
# Installed through `pip install .[dev]`
dev = [
    "pytest",
    "pytest-xdist",
]

# TensorFlow datasets is an extra because it has a large install footprint.
tfds = [
    "tensorflow==2.18.0",
    "tensorflow_datasets==4.9.7",
]

# Grain is an extra because as of v0.2.0 it has no OSX wheels.
grain = [
    "grain==0.2.3",
]

[tool.pyink]
# Formatting configuration to follow Google style-guide
line-length = 80
preview = true
pyink-indentation = 2
pyink-use-majority-quotes = true

[build-system]
build-backend = "setuptools.build_meta"
requires = [
    "setuptools",
]

[tool.setuptools]
packages = ["jax_ai_stack"]
include-package-data = false

[tool.ruff.lint.per-file-ignores]
# F811: Redefinition of unused name.
"docs/source/digits_vae.ipynb" = ["F811"]