pyproject.toml

[project]
name = "jopendde"
version = "0.1.0"
description = "JAX/Equinox translation of OpenDDE"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "dm-tree>=0.1.10",
    "einops>=0.8.0",
    "equinox>=0.11.10",
    "jaxtyping>=0.2.36",
    "numpy>=1.26.3",
]

[project.optional-dependencies]
# Fused NVIDIA cuEquivariance triangle kernels (CUDA 12). Optional: enables the
# triangle-attention kernel via inference.enable_cue_kernels() (a trunk speedup
# that grows with token count). Not imported unless that helper is called.
cue = [
    "cuequivariance-jax>=0.10.0",
    "cuequivariance-ops-jax-cu12>=0.10.0",
]

[dependency-groups]
jax-cpu = [
  "jax[cpu]",
]
jax-cuda = [
  "jax[cuda12]",
]
# The PyTorch OpenDDE reference, vendored under ./OpenDDE. jopendde has no
# featurizer of its own and the weights ship as a torch checkpoint, so
# converting a model / featurizing inputs (Predictor.from_checkpoint,
# Predictor.featurize) need `opendde` + torch. torch/torchvision/torchaudio are
# listed here so the CPU-index source below applies to them.
reference = [
  "opendde",
  "torch",
  "torchvision",
  "torchaudio",
]

[tool.uv.sources]
opendde = { path = "OpenDDE", editable = true }
# torch is only used on CPU (weight extraction); pull the CPU build so it
# doesn't drag in ~5 GB of bundled CUDA wheels. JAX owns the GPU.
torch = { index = "pytorch-cpu" }
torchvision = { index = "pytorch-cpu" }
torchaudio = { index = "pytorch-cpu" }

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[tool.uv]
package = true
conflicts = [
    [
      { group = "jax-cpu" },
      { group = "jax-cuda" },
    ],
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["src/jopendde"]