e3j-ops 0.1.0b2

Fast Euclid equivariant operations for JAX

⚙️ e3j-ops

This package contains the CUDA/C++ source for e3j. The Python bindings to XLA handlers are bundled into the e3j_ops shared object, which the main e3j package wraps in custom JAX primitives via the ffi_call API.

Note: The e3j_ops ABI should not be considered stable for now, but considered private within the e3j Python API.

Building from source

The CMake build recipes are defined in CMakeLists.txt. The Makefile of e3j also defines fine-graind recipes to build and test individual CUDA/C++ objects.

Project structure

The source is organized as follows:

• cuda : CUDA kernel implementations
• ffi : XLA-FFI handlers declarations and pybind11 module definition
• xla : vendored XLA FFI headers
• tests : C++ kernel tests

CUDA/C++ source

Each operation is defined in its own namespace following a common interface (see e.g. tensor_product.cuh for detailed signatures):

• a namespace e3j::op_name within the enclosing e3j namespace,
• a struct e3j::op_name::Params for passing hyperparameters,
• a __global__ CUDA function e3j::op_name::kernel(),
• a __host__ launcher e3j::op_name::launch()

Although object-oriented patterns are hard to mix with __device__ code, this name-based ABI is subject to change.

namespace e3j {
namespace op_name {

    struct Params;

    template <typename Idx, typename Val>
    __global__ void kernel (Params p, ...);

    template <typename Idx, typename Val>
    e3j::Error launch (..., Params p, cudaStream_t stream);

} // namespace op_name
} // namespace e3j

Most kernels are templated across a range of index and/or value data types. Some e3j primitives automatically dispatch to narrow index dtypes (uint8 / uint16) when feature spaces dimensions are small enough. Half-precision float16 arithmetic for values is not yet supported, but planned for a soon upcoming release.

XLA-FFI handlers

The FFI layer uses the XLA FFI API to bind kernel launchers as XLA custom calls. Each handler is registered with XLA_FFI_DEFINE_HANDLER, receiving typed buffers and attributes directly (no opaque descriptor packing):

xla::Error OpNameHandler(
    cudaStream_t stream,
    int32_t num_out,
    xla::AnyBuffer x,
    xla::Result<xla::AnyBuffer> out
) {
    // ... dtype dispatch ...
    return e3j::op_name::launch<Idx, Val>(
        x.typed_data<Idx>(),
        out->typed_data<Val>(),
        params, stream
    ).to_xla();
}

XLA_FFI_DEFINE_HANDLER(
    xla_op_name,
    OpNameHandler,
    xla::Ffi::Bind()
        .Ctx<xla::PlatformStream<cudaStream_t>>()
        .Attr<int32_t>("num_out")
        .Arg<xla::AnyBuffer>()
        .Ret<xla::AnyBuffer>()
);

Handlers are exposed to Python as PyCapsules via pyEncapsulateFunction in e3j_ops.h, and registered in the pybind11 module defined in e3j_ops.cpp.

Contributing

Although it is too early for e3j_ops to accept significant external contributions, bug reports or questions are very welcome via GitHub issues and discussions.

Citing

If you use e3j within your work, we kindly ask you to cite the following preprint:

@article{Peltre26-e3j,
    title   = {{E3J}: an Efficient and Open-Source Euclidean Equivariance Backend},
    author  = {Peltre, Olivier and Picard, Armand and Pichard, Adrien and Giacomoni, Luca and Braganca, Miguel and Heyraud, Valentin and Brunken, Christoph and Tilly, Jules},
    journal = {preprint},
    year    = {2026},
    url     = {(preprint)}
  }
}