fix according to the review

Author: Stellogic

examples/lennard_jones_optimization.py

@@ -79,42 +79,41 @@ def calculate_potential(log_a, epsilon=0.5, sigma=1.0):
 final_a = K.exp(log_a)
 final_energy = calculate_potential(log_a)
 
-if not np.isnan(K.numpy(final_energy)):
-    print("\nOptimization finished!")
-    print(f"Final optimized lattice constant: {final_a:.6f}")
-    print(f"Corresponding minimum total energy: {final_energy:.6f}")
-
-    # Vectorized calculation for the potential curve
-    a_vals = np.linspace(0.8, 1.5, 200)
-    log_a_vals = K.log(K.convert_to_tensor(a_vals))
-
-    # Use vmap to create a vectorized version of the potential function
-    vmap_potential = K.vmap(lambda la: calculate_potential(la))
-    potential_curve = vmap_potential(log_a_vals)
-
-    plt.figure(figsize=(10, 6))
-    plt.plot(a_vals, potential_curve, label="Lennard-Jones Potential", color="blue")
-    plt.scatter(
-        history["a"],
-        history["energy"],
-        color="red",
-        s=20,
-        zorder=5,
-        label="Optimization Steps",
-    )
-    plt.scatter(
-        final_a,
-        final_energy,
-        color="green",
-        s=100,
-        zorder=6,
-        marker="*",
-        label="Final Optimized Point",
-    )
-
-    plt.title("Lennard-Jones Potential Optimization")
-    plt.xlabel("Lattice Constant (a)")
-    plt.ylabel("Total Potential Energy")
-    plt.legend()
-    plt.grid(True)
-    plt.show()
+print("\nOptimization finished!")
+print(f"Final optimized lattice constant: {final_a:.6f}")
+print(f"Corresponding minimum total energy: {final_energy:.6f}")
+
+# Vectorized calculation for the potential curve
+a_vals = np.linspace(0.8, 1.5, 200)
+log_a_vals = K.log(K.convert_to_tensor(a_vals))
+
+# Use vmap to create a vectorized version of the potential function
+vmap_potential = K.vmap(lambda la: calculate_potential(la))
+potential_curve = vmap_potential(log_a_vals)
+
+plt.figure(figsize=(10, 6))
+plt.plot(a_vals, potential_curve, label="Lennard-Jones Potential", color="blue")
+plt.scatter(
+    history["a"],
+    history["energy"],
+    color="red",
+    s=20,
+    zorder=5,
+    label="Optimization Steps",
+)
+plt.scatter(
+    final_a,
+    final_energy,
+    color="green",
+    s=100,
+    zorder=6,
+    marker="*",
+    label="Final Optimized Point",
+)
+
+plt.title("Lennard-Jones Potential Optimization")
+plt.xlabel("Lattice Constant (a)")
+plt.ylabel("Total Potential Energy")
+plt.legend()
+plt.grid(True)
+plt.show()

tensorcircuit/templates/lattice.py

@@ -767,14 +767,11 @@ def _get_distance_matrix_with_mic_vectorized(self) -> Coordinates:
         """
         # Ensure dtype consistency across backends (especially torch) by explicitly
         # casting size and lattice_vectors to the same floating dtype used internally.
-        # Strategy: prefer existing lattice_vectors dtype, fallback to float32 for efficiency,
-        # then float64 for precision. This avoids dtype mismatches in vectorized ops.
-        target_dt = None
-        try:
-            # prefer existing lattice_vectors dtype if possible
-            target_dt = backend.dtype(self.lattice_vectors)  # type: ignore
-        except (AttributeError, TypeError):  # pragma: no cover - defensive
-            target_dt = "float32"
+        # Strategy: prefer existing lattice_vectors dtype; if it's an unusual dtype,
+        # fall back to float32 to avoid mixed-precision issues in vectorized ops.
+        # Note: `self.lattice_vectors` is always created via `backend.convert_to_tensor`
+        # in __init__, so `backend.dtype(...)` is reliable here and doesn't need try/except.
+        target_dt = str(backend.dtype(self.lattice_vectors))  # type: ignore
         if target_dt not in ("float32", "float64"):
             # fallback for unusual dtypes
             target_dt = "float32"
@@ -1626,21 +1623,20 @@ def from_lattice(cls, lattice: "AbstractLattice") -> "CustomizeLattice":
         # Unzip the list of tuples into separate lists of identifiers and coordinates
         _, identifiers, _ = zip(*all_sites_info)
 
-        # Detach-and-copy coordinates in a backend-agnostic way.
-        # Rationale (answering reviewer question "why not keep backend-dependent form?"):
-        # - Passing a tuple/list of backend tensors (e.g., per-row slices) into
-        #   convert_to_tensor can fail on some backends (torch.tensor(list_of_tensors) ValueError),
-        #   whereas a plain nested Python list is accepted everywhere.
-        # - We want CustomizeLattice to be decoupled from the original lattice's computation
-        #   graph and device state (CPU/GPU), so we materialize numeric values here.
-        # - This is a one-shot conversion of the full coordinate array, simpler and faster
-        #   than iterating per row, while preserving the same numeric content.
-        coords_py = backend.numpy(lattice._coordinates).tolist()
+        # Detach-and-copy coordinates while remaining in tensor form to avoid
+        # host roundtrips and device/dtype changes; this keeps CustomizeLattice
+        # decoupled from the original graph but backend-friendly.
+        # Some backends (e.g., NumPy) don't implement stop_gradient; fall back.
+        try:
+            coords_detached = backend.stop_gradient(lattice._coordinates)
+        except NotImplementedError:
+            coords_detached = lattice._coordinates
+        coords_tensor = backend.copy(coords_detached)
 
         return cls(
             dimensionality=lattice.dimensionality,
             identifiers=list(identifiers),
-            coordinates=coords_py,
+            coordinates=coords_tensor,
         )
 
     def add_sites(

tests/test_lattice.py

@@ -2090,14 +2090,12 @@ def test_layering_on_various_lattices(lattice_instance):
     _validate_layers(bonds, layers)
 
 
-# --- Regression tests for backend-scalar lattice constants (PR fix) ---
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
 def test_square_lattice_accepts_backend_scalar_lattice_constant(backend):
     """
     Ensure SquareLattice can be constructed when lattice_constant is a backend scalar tensor
     (e.g., tf.constant, jnp.array, torch.tensor), without mixed-type errors.
     """
-    tc.set_backend(backend)
 
     lc = tc.backend.convert_to_tensor(0.5)
     lat = SquareLattice(size=(2, 2), lattice_constant=lc, pbc=False)
@@ -2121,7 +2119,6 @@ def test_rectangular_lattice_mixed_type_constants(backend):
     RectangularLattice should accept a tuple where one constant is a backend scalar tensor
     and the other is a Python float.
     """
-    tc.set_backend(backend)
 
     ax = tc.backend.convert_to_tensor(0.5)  # tensor scalar
     ay = 2.0  # python float