pyproject.toml

[build-system]
requires = ["flit_core >=3.2,<4"]
build-backend = "flit_core.buildapi"

[project]
name = "optax"
dynamic = ["version"]
description = "A gradient processing and optimisation library in JAX."
readme = "README.md"
license = { file = "LICENSE" }
requires-python = ">=3.9"
authors = [
    {name = "Google DeepMind", email = "optax-dev@google.com"},
]
keywords = [
    "python",
    "machine learning",
    "reinforcement-learning"
]
classifiers = [
    "Environment :: Console",
    "Programming Language :: Python",
    "Intended Audience :: Developers",
    "Operating System :: OS Independent",
    "Programming Language :: Python :: 3",
    "Intended Audience :: Science/Research",
    "Development Status :: 4 - Beta",
    "License :: OSI Approved :: Apache Software License",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
    "Topic :: Software Development :: Libraries :: Python Modules",
]
dependencies = [
  "absl-py>=0.7.1",
  "chex>=0.1.86",
  # Keep jax, jaxlib versions in sync with .github/workflows/tests.yml
  "jax>=0.4.27",
  "jaxlib>=0.4.27",
  "numpy>=1.18.0",
  "etils[epy]",
]

[project.urls]
homepage = "https://github.com/google-deepmind/optax"
repository = "https://github.com/google-deepmind/optax"
documentation = "https://optax.readthedocs.io/"

[project.optional-dependencies]
test = [
    "dm-tree>=0.1.7",
    "flax>=0.5.3",
    "scipy>=1.7.1",
    "scikit-learn"
]

examples = [
  "tensorflow-datasets>=4.2.0",
  "tensorflow>=2.4.0",
  "dp_accounting>=0.4",
  "ipywidgets",
  "flax",
]

docs = [
    "sphinx>=6.0.0",
    "sphinx-book-theme>=1.0.1",  # Older versions fail to pin pydata-sphinx-theme
    "sphinxcontrib-katex",
    "sphinx-autodoc-typehints",
    "ipython>=8.8.0",  # 8.7.0 has ipython3 lexer error
    "myst-nb>=1.0.0",
    "matplotlib>=3.5.0",
    "sphinx-gallery>=0.14.0",
    "sphinx-collections>=0.0.1",
    "tensorflow>=2.4.0",
    "tensorflow-datasets>=4.2.0",
    "flax",
    "sphinx_contributors",
]

dp-accounting = [
    "absl-py>=1.0.0",
    "attrs>=21.4.0",
    "mpmath>=1.2.1",
    "numpy>=1.21.4",
    "scipy>=1.7.1"
]

[tool.setuptools.packages.find]
include = ["README.md", "LICENSE"]
exclude = ["*_test.py"]