pyproject.toml

[project]
name = "tokamax"
authors = [
  {name = "Ani Udipi", email = "udipi@google.com"},
  {name = "Chris Jones", email = "cjfj@google.com"},
  {name = "Christos Perivolaropoulos", email = "cperivol@google.com"},
  {name = "Giorgio Arena", email = "giorgioa@google.com"},
  {name = "Gleb Pobudzey", email = "pobudzey@google.com"},
  {name = "Robert Dyro", email = "rdyro@google.com"},
  {name = "Sebastian Bodenstein", email = "sbodenstein@google.com"},
  {name = "Selim Seferbekov", email = "selimsef@google.com"},
  {name = "Stanley Tzeng", email = "stzeng@google.com"},
  {name = "Niyati Chaudhary", email = "niyatic@google.com"},
]

description = "A Pallas Custom Kernel Library."
readme = "README.md"
license = {file = "LICENSE"}
requires-python = ">=3.11"
dependencies = [
    "absl-py>=2.3.0",
    "einshape",
    "jax>=0.10.2",
    "jaxlib>=0.10.2",
    "jaxtyping>=0.3",
    "pydantic>=2.11.0",
    "qwix>=0.1.2",
    "tqdm",
    "typing_extensions>=4.5.0",
    "typeguard==2.13.3",
    "immutabledict",
    "tensorboardx",
]

dynamic = ["version"]

[tool.setuptools.dynamic]
version = {attr = "tokamax._src.version.TOKAMAX_VERSION"}

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

[tool.setuptools]
include-package-data = true

[tool.setuptools.package-data]
tokamax = ["data/autotuning/*/*.json"]

[project.optional-dependencies]
test = [
    "chex>=0.1.91",
    "flatbuffers",
    "flax",
    "pytest-xdist",
    "pytest",
    "xprof",
    "hypothesis",
]

bench = [
    "google-benchmark>=1.9.0",
    "xprof",
    "cuequivariance-jax>=0.10.0",
    "cuequivariance-ops-jax-cu13>=0.10.0",
    "triton"
]

cuda = [
    "jax[cuda13]>=0.10.1",
    "nvidia-cudnn-cu13>=9.23.0.39",
]

tpu = [
    "jax[tpu]>=0.10.1",
]

[tool.pytest.ini_options]
markers = ["long: tests that will take a long time to run and are used in weekly runs"]