cuEquivariance in PyTorch

Purpose

Execute cuequivariance (cue) descriptors with cuequivariance_torch (cuet): CPU references, CUDA backends, equivariant layers, and autograd. Preserve the project's environment.

Inputs

Required inputs: operation, SO3/O3 irreps, layout, device, dtype, and weight ownership; infer from the request or code. Graphs also need source/destination counts and indices. Optional inputs: backend, performance, export targets.

Instructions

1. Probe once and reuse working PyTorch. Never recursively delete directories during setup or recovery, including partial temporary environments; use a fresh unique path when isolation is needed.
2. Set method="naive" at construction for CPU. Match parameters, features, weights, and references in device and dtype. Use float64 for small numerical and gradient checks.
3. Specify layout and weight ownership explicitly. Derive sizes and output order from the module/descriptor; do not guess multiplicities from a filter.
4. Finish installation before running the saved script with the completion recipe below. A session/task ID, empty output, or a file write is not completion. Poll that session until it exits and inspect its results.
5. For an API mismatch, inspect the installed signature/source once. Adapt or explain the requirement; preserve the requested library.

Prerequisites and environment reuse