JAX API

Added in version v3.0.0.

pySigLib provides a JAX API which exposes all the same functions, but makes them compatible with JAX transformations (jax.jit, jax.grad, jax.vmap). Just import

import pysiglib.jax_api as pysiglib

The JAX API uses XLA FFI to call the same underlying C++/CUDA library, so performance is on par with the base API. All functions support jax.jit compilation and jax.grad differentiation.

Note

pySigLib must be built with JAX FFI support enabled (requires jaxlib >= 0.5.0, Python 3.10+). Check pysiglib.BUILT_WITH_JAX_FFI to verify.

Citation

If you found this library useful in your research, please consider citing the paper:

@article{shmelev2025pysiglib,
  title={pySigLib-Fast Signature-Based Computations on CPU and GPU},
  author={Shmelev, Daniil and Salvi, Cristopher},
  journal={arXiv preprint arXiv:2509.10613},
  year={2025}
}