pyproject.toml

[project]
name = "jaxamples"
version = "0.1.0"
description = "Examples built on JAX, Flax, nnx, onnx, jax2onnx"
authors = [{ name = "enpasos", email = "matthias.unverzagt@enpasos.ai" }]
readme = "README.md"
requires-python = ">=3.12,<4.0"
dynamic = ["dependencies"]

[tool.poetry.dependencies]
python = ">=3.12,<4.0"
# Pinned JAX and related libraries based on working environment
jax = { version = ">=0.7.2", extras = ["cuda12"] }
flax = ">=0.12.1"
ml_dtypes = ">=0.5.1" 
optax = ">=0.2.4"
orbax-checkpoint = ">=0.11.6"
orbax-export = ">=0.0.6"
netron = ">=8.1.9" # Keep as >=, minor updates unlikely to cause segfault
onnx = ">=1.20.1"
onnxruntime = ">=1.24.1"
matplotlib = ">=3.10.0" 
treescope = ">=0.1.9"
jax2onnx = ">=0.10.1"  
# jax2onnx = { path = "../jax2onnx", develop = true }

# Keep PyTorch versions as they are, source correctly points to CPU
torch = { version = "2.5.1", source = "pytorch" }
torchvision = { version = "0.20.1", source = "pytorch" }

[tool.poetry.group.test]
optional = true

[tool.poetry.group.test.dependencies]
matplotlib = "==3.10.0" # Already pinned correctly

[tool.poetry]
packages = [
    { include = "jaxamples" },
]

[[tool.poetry.source]]
name = "jax-releases"
url = "https://storage.googleapis.com/jax-releases/jax_releases.html"
priority = "supplemental"

[[tool.poetry.source]]
name = "pytorch"
url = "https://download.pytorch.org/whl/cpu" # Explicitly CPU source
priority = "explicit"

# [[tool.poetry.source]] # Commenting out testpypi unless needed
# name = "testpypi"
# url = "https://test.pypi.org/simple/"
# priority = "explicit"

[tool.poetry.group.dev.dependencies]
pytest = "^8.3.4" # Keep dev dependencies flexible
pytest-json-report = "^1.5.0"
pre-commit = "^4.1.0"
ruff = "^0.9.5"

[build-system]
requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"

[tool.pytest.ini_options]
testpaths = ["tests"]
pythonpath = ["."]