[pl-fuchs]

This PR is a follow-up to discussion #221 and provides an alternative JAX kernel for UVU convolutions that adapts the receiver-side streaming approach from Vladimir Chorošajev and Cédric Bény.

The main difference from the standard OEQ convolution kernels is the ownership. In the forward pass, each thread owns one (receiver, channel) pair, scans the edges incident to that receiver, and reduces the convolution result locally. In the backward pass, each warp owns one edge, processes the channels in parallel, writes radial-weight gradients directly, reduces spherical-harmonic gradients across channels within the warp, and performs atomic additions for node feature gradients. Unlike the SOBEK schedule, the implementation consumes materialized per-edge weights, is specialized to UVU convolutions, and does not support general UVW convolutions.

The PR moves the standard JAX Oeq convolutions and the new streaming convolution behind a common public interface. Either implementation can be explicitly selected via a mode argument, if supported. The auto mode selects the streaming path for UVU and reducible UVW convolutions and falls back to the standard path otherwise.

I developed this PR with assistance from Codex. I have read and reviewed the submitted changes. I tested the changes on CUDA, but have not tested the HIP backend.

[vbharadwaj-bk]

Thanks for submitting this! Bigger diff than I expected, will start looking at it this week (last week was a push to get @abagusetty 's XPU support merged).

[pl-fuchs]

Thank you for the effort. If you have an idea how I could align it more tightly with the existing code, I will give it a try.

[asglover]

"The PR moves the standard JAX Oeq convolutions and the new streaming convolution behind a common public interface."

I'm not sure, I'd like to do this, I think I'd rather keep them separate as they literally have different function signatures. I would rather applications make the choice to develop wrappers which implement layers by calling either implementation. I'll give this area a close read when I get the chance.

[pl-fuchs]

Both implementations take the same core inputs (X, Y, W, rows, cols) and return updated node features. Streaming supports a subset of tensor-product configurations, while the extra arguments are mode-specific, similar to sender_perm in deterministic mode.