pyproject.toml

[build-system]
requires = ["setuptools"]
build-backend = "setuptools.build_meta"

[project]
name = "pyjuice"
version="2.7.0"
description = "Probabilistic Circuits Library"
dependencies = [
    "numpy",
    "torch>=2.2.0",
    "typing",
    "triton>=2.2.0",
    "networkx",
    "numba",
    "packaging",
    "matplotlib",
    "tqdm"
]
authors = [
    {name="Anji Liu", email="anjiliu@comp.nus.edu.sg"},
    {name="StarAI", email="guyvdb@cs.ucla.edu"},
]

[project.optional-dependencies]
# Optional CuTe/CUTLASS + TMA sum-layer fast-path kernels (sm_90+; falls back to Triton otherwise).
# `pip install pyjuice[cuda]` ships the header-only CUTLASS so the fast path can JIT-compile; a plain
# install stays Triton-only and needs no CUDA toolchain. See pyjuice/layer/kernels/c/__init__.py.
cuda = [
    "nvidia-cutlass>=3.5",
]
# (kept empty: text-level constraints need nothing beyond the base install since they read Python's own regex parse)
constraints = []
dev = [
    "pytest",
    "pytest-xdist",
    "pytest-skip-slow",
    "torchvision",
    "datasets",
    "matplotlib",
    "jsonschema",
    "outlines-core>=0.2.14,<0.3",      # (tests: the JSON Schema texts against outlines-core's)
]

[options.packages.find]
where = "src"

[tool.setuptools.dynamic]
readme = {file = "README.md"}

# Ship the CUDA source for the optional JIT-compiled fast-path kernels with the package, so it is
# present in installed trees / wheels (it is compiled at runtime via torch.utils.cpp_extension).
[tool.setuptools.package-data]
"pyjuice.layer.kernels.c" = ["*.cu", "*.cuh", "*.h"]
# Nested kernel packages also need their sources and headers for runtime compilation.
"pyjuice.layer.kernels.c.input_layers" = ["*.cu"]
"pyjuice.nodes.distributions.c_kernels" = ["*.cu"]
"pyjuice.nodes.external_params.kernels.c" = ["*.cu", "*.cuh"]
"pyjuice.constraints.inference.dfa_exact.kernels.c" = ["*.cu", "*.cuh"]
"pyjuice.utils.c" = ["*.cpp"]

[tool.pytest.ini_options]
addopts = [
    "--import-mode=importlib",
]
testpaths = [
    "tests/"
]