BCSR support for mul

Author: refraction-ray

examples/analog_rydberg.py

@@ -21,7 +21,9 @@
 
 # Instantiate the AnalogCircuit
 ac = tc.AnalogCircuit(nqubits)
-
+ac.set_solver_options(
+    ode_backend="diffrax", max_steps=20000
+)  # more efficient and stable than the default one by jax
 # 1. Create the sparsely excited state
 ac.x([i for i in range(nqubits) if i % 4 == 0])
 
@@ -49,5 +51,7 @@ def rydberg_hamiltonian_func(t):
 #         ac.rx(i, theta=-np.pi/4)
 
 # 5. Sample from the final state in the computational basis
+state = ac.state()
+np.testing.assert_allclose(tc.backend.norm(state), 1, atol=1e-3)
 sample = ac.sample(batch=1024, allow_state=True, format="count_dict_bin")
 print("\nSampled bitstrings:\n", sample)

tensorcircuit/backends/jax_backend.py

@@ -175,6 +175,27 @@ def _eigh_jax(self: Any, tensor: Tensor) -> Tensor:
     return adaware_eigh(tensor)
 
 
+def bcsr_scalar_mul(self: Tensor, other: Tensor) -> Tensor:
+    """
+    Implements scalar multiplication for BCSR matrices (self * scalar).
+    """
+    import jax.numpy as jnp
+    from jax.experimental.sparse import BCSR
+
+    if jnp.isscalar(other):
+        # The core logic: only the data array is affected by scalar multiplication.
+        # The sparsity pattern (indices, indptr) remains the same.
+        new_data = self.data * other
+
+        # Return a new BCSR instance with the scaled data.
+        return BCSR((new_data, self.indices, self.indptr), shape=self.shape)
+
+    # For any other type of multiplication (e.g., element-wise with another matrix),
+    # return NotImplemented. This allows Python to try other operations,
+    # like other.__rmul__(self).
+    return NotImplemented
+
+
 tensornetwork.backends.jax.jax_backend.JaxBackend.convert_to_tensor = (
     _convert_to_tensor_jax
 )
@@ -224,6 +245,11 @@ def __init__(self) -> None:
 
         self.name = "jax"
 
+        # --- Monkey-patch the BCSR class ---
+
+        sparse.BCSR.__mul__ = bcsr_scalar_mul  # type: ignore
+        sparse.BCSR.__rmul__ = bcsr_scalar_mul  # type: ignore
+
     # it is already child of numpy backend, and self.np = self.jax.np
     def eye(
         self, N: int, dtype: Optional[str] = None, M: Optional[int] = None

tensorcircuit/backends/pytorch_backend.py

@@ -24,7 +24,6 @@
 
 # TODO(@refraction-ray): lack stateful random methods implementation for now
 # TODO(@refraction-ray): lack scatter impl for now
-# TODO(@refraction-ray): lack sparse relevant methods for now
 # To be added once pytorch backend is ready
 
 

tensorcircuit/timeevol.py

@@ -616,7 +616,7 @@ def ode_evol_global(
 
     def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
         h = -1.0j * hamiltonian(t, *args)
-        return backend.sparse_dense_matmul(h, y)
+        return h @ y
 
     s1 = _solve_ode(f, initial_state, times, args, solver_kws)