Fix JAX tensor product backward buffer initialization

[pl-fuchs]

Thanks for building and maintaining OpenEquivariance. While applying the JAX tensor product path to MACE, I ran into a gradient correctness issue with FCTP-style tensor products.

Summary

This PR fixes a JAX TensorProduct gradient correctness issue in the FFI path. The
backward and double-backward kernels now initialize their gradient outputs before
accumulating into them.

Problem

The JAX FFI result buffers returned by tp_backward and tp_double_backward
were treated as if they started at zero. JAX does not guarantee that for FFI
outputs. Since the OEQ backward kernels accumulate into these buffers, any
pre-existing contents in the uninitialized result buffers could be added to
otherwise correct gradients. This showed up when multiple JAX tensor product
operators were constructed and used in the same process, especially in
shared-weight/FCTP-style MACE usage.

The PR also adds a defensive TP cache guard. The compiled kernel cache now
stores the serialized kernel JSON alongside the truncated SHA cache key and
raises if a later call presents the same key with a different payload.
The collision path is not believed to be the cause of the reported gradient bug, but
it prevents a future or manually triggered key collision from silently reusing a
kernel compiled for a different payload.

Fix

This PR zeroes all gradient result buffers before launching the JAX TP backward
and double-backward kernels.

The regression test constructs multiple JAX TP operators in one process and
checks JVP/VJP adjoint consistency. It covers shared uvu and unshared uvw
cases, and it exercises both the backward and double-backward FFI paths.

Tests

pre-commit run --all-files
OEQ_NOTORCH=1 python -m pytest tests/tp_adjoint_test.py --jax -q -rs
python -m pytest tests/import_test.py -q -rs

[vbharadwaj-bk]

Thanks much for fixing this! I would request that the kernel cache collision checking be reverted (see below for an explanation and an alternate way to handle that avoids the expensive string comparison) and kept for a separate PR. We can get this merged summarily.

[vbharadwaj-bk]

Hmm - this line is a bit problematic since we call this compile_tp_with_caching before every kernel invocation, relying on cache hits to short circuit expensive computation. The payload strings can be several hundred kilobytes and increases runtime by precious microseconds to do the string comparison.

That said, the concern about hash collisions is valid. One way to get around this (probably best left for another PR to keep things clean) would be to add a "nonce" integer field to the json with a default value of zero, along with a function called "check_for_collision" called from Python. TP initialization would call this function before compilation and look up whether a) the hash exists and b) would compare the payload json for equality; the expensive check is fine here because it's only done once at TP construction time. In case we do have a hash collision, we could just tick up the nonce in the json until no collision occurs, then proceed with our current flow. @asglover, in case you are interested in implementing this.

[asglover]

What do you think about flipping the API design, so that the C++ library choses and returns the ID? Because it has access to the hash table, it can repeatedly probe the table for empty spots, then it could return a unique ID to the caller?

[vbharadwaj-bk]

Hmm - to my knowledge, black-box hash tables like std::map don't usually expose their internals in this manner; if we want a lookup key that's based on a hash of the input payload, the only method I can think of is modifying the payload itself (e.g. using the nonce), which, for consistency means reserializing and storing the json to the payload tensor (I'd prefer not to modify the payload at both the C++ and python layers, not to mention that the deserialized json might be a different byte length than the original).

The alternative is to hand-roll a hashtable that returns a unique slot - but even std hash tables are fairly optimized, and this seems like a bit of code to write.

Either way, eliminating this problem (which we haven't run into yet) still requires splitting into a slow string-checking precompile function and the fast lookup function, both exposed at the Python layer. I'm also not sure how much effort we want to put into solving this given that our users compile very few unique TPs (I estimate ~billions of tensor products for a 50% collision chance with a 63-bit key). The simplest fix of all would be doubling the hash key length to 128 bits, which obliterates failure probability.

[asglover]

There is an API for checking this in std::unordered map bucket(key). Although, getting the no collision lookup is really only a performance factor. I probably shouldn't worry about it.

I was proposing having the lookup key not be related to a hash of the input payload. The api would be that the user provides the input payload and then the TP_JIT returns the key.

The only correctness invariant is that two different input payloads never are associated with the same key.

The TP-jit could actually just mint new IDs in order from zero. It is the owner of all the implementations. And the lookup could a vector access.

One thing I have been thinking about is: How critical is it that the JIT library identify when someone request the same payload twice - and skip compilation. I was thinking that it would be nice to not promise this behavior and semantically always return a new implementation. But I don't know if someone is really relying on this behavior and minting 1000's of identical TPs. It feels like a more transparent contract to always provide a new implementation and the user can always build sharing into the their model builder implementation by explicitly caching and reusing a JIT key, or custom op library namespace. If you feel strongly about this please push back. In practice, the implementation could return the same key if it promises to be correct about it.

I was leaning towards UUID. However, it is nice that FFI has types for uint64. a 128 bit ID would have to be passed in a tensor or by pointer.

I think in this instance, if you get the keys from the JIT library, you have the ability to check and adjust your IDs with the only registry that matters (the JIT library's registry), so you can never have conflicting keys. So I think doing a uint64 is totally reasonable. You would literally need to use all 2^64 slots to run out.

I was thinking about what is needed for doing a one-TP-is-one-torch-cpp-extension API. You can actually check for name collisions when you're registering the extension (and guarantee no collisions). But when the user exports their model, and they go load it up in some C++ based molecular simulation runtime, the names really need to be random and in a large space. If they collided, I think torch would complain at load time, which is not catastrophic, but still annoying. So that had me leaning towards randomly assigning a UUID. And there's no "passing by pointer issues" because the UUID would just be baked into the name of the library. Something like "openequivariance_XXXX_XXXX_XXXX_XXXX::TensorProductForward"

let me know your thoughts

[vbharadwaj-bk]

I see, I didn't know about the bucket function! Cool.

With that said - I prefer the key be a hash of the input, and not to not change the existing kernel cache mechanism; to me, this is optimal and a nice shared mechanism that is similar across both PyTorch and JAX.

Minting new IDs from zero, for example, would be a problem for torch.compiled models; suppose that someone loads up two distinct compiled models at the C++ layer (very realistic in an inference setting). If each model was compiled in isolation, kernel 0 could be very different between the two. For torch.compiled or JAX compiled models, we always rely on the fast hash lookup and ideally, wouldn't do any payload string comparison, and therefore a kernel ID of zero for both would be a problem. That could happen with a hash as well (hash collision) - but at least the probability is minimized. As for bucket ID: this is dependent on the hashtable length - makes me a bit wary.

[asglover]

With that said - I prefer the key be a hash of the input, and not to not change the existing kernel cache mechanism; to me, this is optimal and a nice shared mechanism that is similar across both PyTorch and JAX.

Yeah, the shared mechanism is nice. I've been thinking about trying to move the jinja templating and greedy algorithm into C++ so that the api would be TP-JIT(TPP) -> key. Then the bindings would be really clean. All the implementation logic would be behind the C++ bindings.

Minting new IDs from zero, for example, would be a problem for torch.compiled models; suppose that someone loads up two distinct compiled models at the C++ layer (very realistic in an inference setting). If each model was compiled in isolation, kernel 0 could be very different between the two. For torch.compiled or JAX compiled models, we always rely on the fast hash lookup and ideally, wouldn't do any payload string comparison, and therefore a kernel ID of zero for both would be a problem. That could happen with a hash as well (hash collision) - but at least the probability is minimized.

I'm not sure about this analysis. I actually think it might be the opposite. When torch.compile exports for C++ inference, the actual JIT library dies, and retains no state. The kernels persist as u8 encoded strings in a module buffer tensor. Then when the c++ inference starts up, the kernel string are fed into the new TP_jit library. I think that the only reason that a hash collision could exist at this point in time is because the key does not get to be chosen by the TP-JIT library, which owns the registry.

If you have some situation where there are two different kernel strings that both deterministically hash to the same key, then you have a collision and it's a correctness problem.

If the JIT-TP library has an internal mutex, and can only register kernels sequentially, it will receive those kernel strings sequentially and can assign them non-colliding keys.

I think with the JIT-TP library mechanism, because the only registration that matters is inside the library, if the library choses the keys, you can actually guarantee no collision. And you can do this with the C++ extension mechanism too, for a single process, (by checking before registering a name). The only time that you can't guarantee no collisions is when you export C++ extensions into some other process. That's why the C++ extension names need to be random and 128 bits of random for sanity. But I think the JIT-TP ids can always be collision free and so can safely just be 64 bits.

Is my argument making sense? Or am I forgetting a key detail?
The design is slightly more difficult if you also want to guarantee that the same input string must be a collision and reuse.
But I don't think this is something we should guarantee.
And even if it was, it could be implemented in C++ and retain the correctness guarantee of never having a false collision.

As for bucket ID: this is dependent on the hashtable length - makes me a bit wary.

Having the same bucket ID doesn't actually represent a collision. So long as the keys are different, each will correctly point to the correct value. It's just that the accessor will need to do a linear scan within a bucket, so it would incur a slight performance hit.

[asglover]

also just to clarify, I don't think this is a big risk and am happy to do it whichever way

[vbharadwaj-bk]

Awesome, thanks!

[pl-fuchs]

Thanks for the quick review and explanation. I reverted the kernel cache collision check and kept this PR scoped to the JAX FFI buffer initialization fix only.

[vbharadwaj-bk]

✅