Metadata-Version: 2.4
Name: openequivariance_extjax
Version: 0.7.0
Summary: JAX C++ Extension for OpenEquivariance
Author: Austin Glover, Vivek Bharadwaj, Aydin Buluc, James Demmel
License-Expression: BSD-3-Clause
License-File: LICENSE
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Project-URL: homepage, https://passionlab.github.io/OpenEquivariance/
Project-URL: source, https://github.com/PASSIONLab/OpenEquivariance
Project-URL: issues, https://github.com/PASSIONLab/OpenEquivariance/issues
Requires-Python: >=3.10
Description-Content-Type: text/markdown

# OpenEquivariance JAX Extension

The JAX extension module for OpenEquivariance.

## XLA FFI provider

The `libjcn_ffi_openequivariance.so` Bazel target builds a Python-independent
FFI provider against the XLA and CUDA repositories selected by the caller's
Bazel configuration:

```text
bazel build @openequivariance_src//openequivariance_extjax:libjcn_ffi_openequivariance.so
```

The handler implementation owns one ABI-versioned table with slots for every
XLA stage (`instantiate`, `prepare`, `initialize`, and `execute`) plus each
target's traits. The nanobind extension reads that table and publishes only
stages that have a handler. The shared library exports `RegisterFFi(const
XLA_FFI_Api*, const char*)`; a runtime loader obtains the API from its PJRT
plugin and calls this entry point. The provider uses
`XLA_FFI_Handler_Register` to register every target as a complete handler
bundle, passing every non-null stage and the traits from the table. JAX
registration uses the same table.
