pyproject.toml

[project]
name = "jax-flow-matching"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = "~=3.11"
dependencies = [
    "clu>=0.0.12",
    "datasets>=3.0.1",
    "diffrax>=0.6.0",
    "flax>=0.10.0",
    "huggingface-hub>=0.26.0",
    "jax>=0.4.34 ; sys_platform == 'darwin'",
    "jax[cuda]>=0.4.34 ; sys_platform == 'linux'",
    "jaxtyping>=0.2.34",
    "matplotlib>=3.9.2",
    "numpy<2",
    "tensorflow-cpu>=2.14.0",
    "pydantic>=2.9.2",
    "wandb>=0.18.5",
    "orbax-checkpoint",
]

[tool.ruff.lint]
select = ["E", "F", "I001"]
ignore = ["E721", "E731", "F722", "F821"]

[tool.mypy]
disable_error_code = "import-untyped"

[tool.uv]
dev-dependencies = ["mypy>=1.12.0", "ruff>=0.7.0"]
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]

[tool.uv.sources]
orbax-checkpoint = { git = "https://github.com/google/orbax/", subdirectory = "checkpoint" }