JVP Support (#181)

* Added a JSN parser.

* Debugging in progress.

* Bound the double backward pass.

* Ready to try conv next.

* Adding the convolution primitives.

* Convolution working.

* JITted.

* Tests JIT also.

* Third derivative is failing.

* Fixed failing third derivative.

* Fixed up some more stuff.

* Added VJP version for debugging.

* Fixed zero buffers.

* Zero'd some more buffers.

* More things are working.

* Fixed things up.

* Precommit.

* Merge changes.

Author: vbharadwaj-bk

CHANGELOG.md

@@ -1,5 +1,22 @@
 ## Latest Changes
 
+### v0.5.4 (2025-02-01) 
+Improvements to JAX frontend.
+
+**Added**:
+- Jacobian Vector Products (JVP) 
+  for both `TensorProduct` and `TensorProductConv` via custom primitives, in addition to VJP.
+- Arbitrary higher-order derivatives in JAX.
+- JAX JIT support; in particular, support for
+  Phonon Fine Tuning in [Nequix](https://github.com/atomicarchitects/nequix).
+
+**Fixed**:
+- Zero'd all output buffers in the backwards and double-backwards implementations of convolution
+before calling kernels. 
+
+### v0.5.1-0.5.3 (2025-02-01) 
+Minor bugfixes related to packaging and JAX. 
+
 ### v0.5.0 (2025-12-25) 
 JAX support is now available in 
 OpenEquivariance for BOTH NVIDIA and

openequivariance/openequivariance/core/ComputationSchedule.py

@@ -320,7 +320,7 @@ def __init__(self, num_blocks, num_threads, warp_size, smem):
         self.num_blocks = num_blocks
         self.num_threads = num_threads
         self.warp_size = warp_size
-        self.smem = smem
+        self.smem = int(smem)
 
 
 class ComputationSchedule:

openequivariance/openequivariance/core/utils.py

@@ -7,7 +7,6 @@
 
 import json
 import tempfile
-import hashlib
 
 from enum import IntEnum
 
@@ -200,13 +199,3 @@ def benchmark(func, num_warmup, num_iter, mode="gpu_time", kernel_names=[]):
             time_millis[i] = kernel_time
 
     return time_millis
-
-
-def hash_attributes(attrs):
-    m = hashlib.sha256()
-
-    for key in sorted(attrs.keys()):
-        m.update(attrs[key].__repr__().encode("utf-8"))
-
-    hash = int(m.hexdigest()[:16], 16) >> 1
-    attrs["hash"] = hash