fix scatter in jax backend

Author: refraction-ray

tensorcircuit/backends/jax_backend.py

@@ -609,16 +609,16 @@ def f_jax(*args: Any, **kws: Any) -> Any:
         return carry
 
     def scatter(self, operand: Tensor, indices: Tensor, updates: Tensor) -> Tensor:
-        updates = jnp.reshape(updates, indices.shape)
-        return operand.at[indices].set(updates)
-        # rank = len(operand.shape)
-        # dnums = libjax.lax.ScatterDimensionNumbers(
-        #     update_window_dims=(),
-        #     inserted_window_dims=tuple([i for i in range(rank)]),
-        #     scatter_dims_to_operand_dims=tuple([i for i in range(rank)]),
-        # )
-        # r = libjax.lax.scatter(operand, indices, updates, dnums)
-        # return r
+        # updates = jnp.reshape(updates, indices.shape)
+        # return operand.at[indices].set(updates)
+        rank = len(operand.shape)
+        dnums = libjax.lax.ScatterDimensionNumbers(
+            update_window_dims=(),
+            inserted_window_dims=tuple([i for i in range(rank)]),
+            scatter_dims_to_operand_dims=tuple([i for i in range(rank)]),
+        )
+        r = libjax.lax.scatter(operand, indices, updates, dnums)
+        return r
 
     def coo_sparse_matrix(
         self, indices: Tensor, values: Tensor, shape: Tensor

tensorcircuit/quantum.py

@@ -1161,6 +1161,9 @@ def quimb2qop(qb_mpo: Any) -> QuOperator:
     return qop
 
 
+# TODO(@refraction-ray): Z2 analogy or more general analogies for the following u1 functions
+
+
 def u1_inds(n: int, m: int) -> Tensor:
     """
     Generate all the combination index of m down spins in n sites.