fix

e3nn · Sep 13, 2023 · bc1b921 · bc1b921
1 parent 06142e9
commit bc1b921
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/e3nn_jax/_src/irreps_array.py b/e3nn_jax/_src/irreps_array.py
@@ -1118,6 +1118,7 @@ def rechunk(self, irreps: IntoIrreps) -> "IrrepsArray":
 
         new_chunks = None
         if self._chunks is not None:
+            jnp = _infer_backend(self.array)
             leading_shape = self.shape[:-1]
 
             new_chunks = []