From 9853e86deb76971ad5973ad1d1f77153c7e9eda9 Mon Sep 17 00:00:00 2001
From: Mattia Robbiano <mattiarobbiano@gmail.com>
Date: Sat, 20 Sep 2025 15:15:40 +0200
Subject: [PATCH] fixing grad nan

---
 examples/quimb_intro/test.py |  6 +--
 src/qibotn/backends/quimb.py | 78 +++++++++++++++++++++++++++++++-----
 2 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/examples/quimb_intro/test.py b/examples/quimb_intro/test.py
index a6f1d7e..da499f2 100644
--- a/examples/quimb_intro/test.py
+++ b/examples/quimb_intro/test.py
@@ -30,13 +30,14 @@ def build_circuit(nqubits, nlayers):
         for q in range(nqubits):
             circ.add(gates.RY(q=q, theta=0.))
             circ.add(gates.RZ(q=q, theta=0.))
-        [circ.add(gates.CNOT(q%nqubits, (q+1)%nqubits) for q in range(nqubits))]
+        [circ.add(gates.CZ(q % nqubits, (q + 1) % nqubits)) for q in range(nqubits)]
     circ.add(gates.M(*range(nqubits)))
     return circ
 
 nqubits = 4
 circuit = build_circuit(nqubits=nqubits, nlayers=3)
 
+quimb_circuit = quimb_backend._qibo_circuit_to_quimb(circuit)
 
 def f(params):
     circuit.set_parameters(params)
@@ -46,5 +47,4 @@ def f(params):
     )
 
 parameters = np.random.uniform(-np.pi, np.pi, size=len(circuit.get_parameters()))
-print(f(parameters))
-print(jax.value_and_grad(f)(parameters))
\ No newline at end of file
+print(jax.value_and_grad(f)(parameters))
diff --git a/src/qibotn/backends/quimb.py b/src/qibotn/backends/quimb.py
index b25e559..3f09859 100644
--- a/src/qibotn/backends/quimb.py
+++ b/src/qibotn/backends/quimb.py
@@ -14,27 +14,35 @@ from qibo.result import QuantumState
 from qibotn.backends.abstract import QibotnBackend
 from qibotn.result import TensorNetworkResult
 
+from qibo.gates.abstract import ParametrizedGate
+
+
 GATE_MAP = {
         "h": "H",
         "x": "X",
         "y": "Y",
         "z": "Z",
         "s": "S",
+
         "sdg": "SDG",
         "t": "T",
         "tdg": "TDG",
         "sx": "SX",
         "sxdg": "SXDG",
+
         "rx": "RX",
         "ry": "RY",
         "rz": "RZ",
+
         "u1": "U1",
         "u2": "U2",
         "u3": "U3",
-        "cx": "CNOT",
+
+        "cx": "CX",
         "cnot": "CNOT",
         "cy": "CY",
         "cz": "CZ",
+
         "iswap": "ISWAP",
         "swap": "SWAP",
         "ccx": "CCX",
@@ -368,6 +376,45 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
 
         return A_new, B_new, C_new
 
+    # def _qibo_circuit_to_quimb(self, qibo_circ, quimb_circuit_type=qtn.Circuit, **circuit_kwargs):
+    #     """
+    #     Convert a Qibo Circuit to a Quimb Circuit.
+
+    #     Parameters
+    #     ----------
+    #     qibo_circ : qibo.models.circuit.Circuit
+    #         The circuit to convert.
+    #     quimb_circuit_type : type
+    #         The Quimb circuit class to use (Circuit, CircuitMPS, etc).
+    #     circuit_kwargs : dict
+    #         Extra arguments to pass to the Quimb circuit constructor.
+
+    #     Returns
+    #     -------
+    #     circ : quimb.tensor.circuit.Circuit
+    #         The converted circuit.
+    #     """
+    #     nqubits = qibo_circ.nqubits
+    #     quimb_gates = []
+    #     circ = quimb_circuit_type(nqubits, **circuit_kwargs)
+
+    #     for gate in qibo_circ.queue:
+    #         gname = getattr(gate, "name", None)
+    #         qname = GATE_MAP.get(gname, None)
+    #         if qname is None:
+    #             continue  # skip measurements and unknown gates
+
+    #         # Handle parametrized gates (Qibo: .parameters, Quimb: expects flat tuple)
+    #         params = getattr(gate, "parameters", ())
+    #         qubits = getattr(gate, "qubits", ())
+
+    #         # Quimb expects (*params, *qubits)
+    #         gate_spec = (qname,) + tuple(params) + tuple(qubits)
+    #         quimb_gates.append(gate_spec)
+
+    #     circ.apply_gates(quimb_gates)
+    #     return circ
+
     def _qibo_circuit_to_quimb(self, qibo_circ, quimb_circuit_type=qtn.Circuit, **circuit_kwargs):
         """
         Convert a Qibo Circuit to a Quimb Circuit.
@@ -387,7 +434,7 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
             The converted circuit.
         """
         nqubits = qibo_circ.nqubits
-        quimb_gates = []
+        circ = quimb_circuit_type(nqubits, **circuit_kwargs)
 
         for gate in qibo_circ.queue:
             gname = getattr(gate, "name", None)
@@ -395,14 +442,23 @@ class QuimbBackend(QibotnBackend, NumpyBackend):
             if qname is None:
                 continue  # skip measurements and unknown gates
 
-            # Handle parametrized gates (Qibo: .parameters, Quimb: expects flat tuple)
             params = getattr(gate, "parameters", ())
             qubits = getattr(gate, "qubits", ())
-
-            # Quimb expects (*params, *qubits)
-            gate_spec = (qname,) + tuple(params) + tuple(qubits)
-            quimb_gates.append(gate_spec)
-
-        circ = quimb_circuit_type(nqubits, **circuit_kwargs)
-        circ.apply_gates(quimb_gates)
-        return circ
\ No newline at end of file
+            
+            # Check if the gate is parametrized
+            is_parametrized = isinstance(gate, ParametrizedGate)
+            
+            if is_parametrized:
+                circ.apply_gate(
+                    qname,
+                    *params,
+                    *qubits,
+                    parametrized= is_parametrized
+                )
+            else:
+                circ.apply_gate(
+                    qname,
+                    *params,
+                    *qubits,
+                                )
+        return circ