Skip to content

ML-KULeuven/klay

Repository files navigation

KLay

Python 3.10 License PyPI Downloads

KLay is a Python library for evaluating sparse circuits on the GPU.

To get started, install KLay using pip and check out the documentation. You can also refer to this video or the paper for more information.

pip install klaycircuits

KLay features:

  • GPU acceleration of circuits using Jax or PyTorch.
  • Loading SDD and d-DNNF circuits compiled by PySDD or D4.
  • Evaluation in various semirings (e.g. real, log, tropical).
  • Propagating constants and merging duplicate nodes.

🧪 Tests

Run the test suite from the project root:

pytest tests/

Tests are split by backend. They are automatically skipped if the required backend is not installed:

  • tests/test_manual.py, tests/fuzzer_torch.py, tests/fuzzer_torch_multi.py, tests/fuzzer_creation.py — require PyTorch
  • tests/fuzzer_jax.py — requires JAX
  • tests/test_compression.py — no backend required

📊 Benchmarks

Benchmarks live in the benchmark/ directory. Run from the project root with:

python -m benchmark.benchmark_wmc --benchmark sdd --target torch -v 100 200 500
python -m benchmark.benchmark_wmc --benchmark sdd --target jax   -v 100 200 500
python -m benchmark.benchmark_wmc --benchmark sdd --target pysdd -v 100 200 500

Key options:

Flag Description
-b / --benchmark Circuit type: sdd or d4
-t / --target Backend: torch, jax, or pysdd
-v / --nb_vars Number of variables (one or more)
-d / --device Device: cpu, cuda, cuda:0, etc.
-s / --semiring Semiring: log (default) or real
-r / --nb_repeats Number of seeds to average over (default: 1)

Results are saved as JSON files under results/.

📃 Paper

If you use KLay in your research, consider citing our paper.

To replicate the exact results and figures of the paper, use this code.

@inproceedings{
    maene2025klay,
    title={{KL}ay: Accelerating Arithmetic Circuits for Neurosymbolic {AI}},
    author = {Maene, Jaron and Derkinderen, Vincent and Zuidberg Dos Martires, Pedro},
    booktitle={The Thirteenth International Conference on Learning Representations},
    year={2025},
    url={https://openreview.net/forum?id=Zes7Wyif8G}
}