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.

[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.