cuEquivariance in JAX

Purpose

Execute cuequivariance (cue) descriptors with cuequivariance_jax (cuex). Preserve the project's framework and devices. Custom representations use core cuequivariance; PyTorch modules use cuequivariance_torch.

Inputs

Required inputs: operation, SO3/O3 irreps, layout, batch shapes, dtype, device, and weight ownership. Derive them from the request/code; state defaults. Graphs also need indices and destination count. Operand sizes/output irreps come from descriptors. Optional inputs: backend/performance preferences.

Instructions

1. Probe the project interpreter once and reuse working packages. Environment setup/recovery must not recursively delete directories, including temporary environments. Use a fresh unique path when isolation is needed.
2. Select method="naive" for CPU; Flax is needed only for NNX. High-level dictionary blocks are (..., mul, ir.dim); uniform CUDA arrays differ.
3. Adapt one matching recipe into a script. Build descriptors outside transformed functions, broadcast leading input axes, and enable x64 before array creation for float64 checks.
4. Finish installation before running the saved script with the completion recipe below. A running session/task ID, empty output, or a file write is not a completed check. Poll the same session until it exits; then read its results before reporting success.
5. Investigate a concrete failure using its error and the affected installed API. Rerun affected checks after a fix. Once the requested checks pass, finish; avoid repeated probes, unrelated examples, or exploration of other backends for a CPU reference.

Prerequisites and environment recovery