pyproject.toml

[project]
name = "torch_jax_interop"
description = "Simple tools to mix and match PyTorch and Jax - Get the best of both worlds!"
readme = "README.md"
authors = [
    { name = "Fabrice Normandin", email = "fabrice.normandin@gmail.com" },
]
requires-python = ">=3.11"
dependencies = [
    "jax>=0.6.0",
    "torch",
    # note: This is because of a weird bug where torch wants setuptools to build cpp extensions (seems related to torch.compile).
    "setuptools; python_version == '3.11'",
]
dynamic = ["version"]

[project.optional-dependencies]
gpu = ["jax[cuda13]; sys_platform == 'linux'"]


[dependency-groups]
dev = [
    "pytest>=8.3.3",
    "pytest-cov>=5.0.0",
    "uv-dynamic-versioning>=0.2.0",
    "mkdocs-material>=9.5.44",
    "flax>=0.10.2",
    "tensor-regression>=0.0.8",
    "mktestdocs>=0.2.4",
    "pytest-benchmark>=5.1.0",
    "pytest-env>=1.1.5",
    "mkdocstrings[python]>=0.27.0",
    "black>=24.10.0",
]

[tool.pytest.ini_options]
testpaths = ["torch_jax_interop"]
norecursedirs = [".venv"]
addopts = ["--doctest-modules"]

[tool.pytest_env]
CUBLAS_WORKSPACE_CONFIG = ":4096:8"

[tool.ruff]
line-length = 99

[tool.docformatter]
wrap-summaries = 99
wrap-descriptions = 99

[tool.uv]
managed = true

[tool.uv-dynamic-versioning]
vcs = "git"
style = "semver"

[build-system]
requires = ["hatchling", "uv-dynamic-versioning"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["torch_jax_interop"]

[tool.hatch.version]
source = "uv-dynamic-versioning"