fix according to the review

Author: Stellogic

examples/lennard_jones_optimization.py

@@ -12,16 +12,7 @@
 import optax
 import numpy as np
 import matplotlib.pyplot as plt
-
-# Try to enable JAX 64-bit precision if available (safe fallback)
-
-try:  # pragma: no cover - optional optimization
-    from jax import config as jax_config  # type: ignore
-
-    jax_config.update("jax_enable_x64", True)
-except Exception:  # broad: environment may not have config attribute
-    pass
-import tensorcircuit as tc  # noqa: E402
+import tensorcircuit as tc
 
 
 tc.set_dtype("float64")  # Use tc for universal control
@@ -58,9 +49,8 @@ def calculate_potential(log_a, epsilon=0.5, sigma=1.0):
     return potential_energy
 
 
-# Create a lambda function for optimization
-potential_fun_for_grad = lambda log_a: calculate_potential(log_a)
-value_and_grad_fun = K.jit(K.value_and_grad(potential_fun_for_grad))
+# Create value and grad function for optimization
+value_and_grad_fun = K.jit(K.value_and_grad(calculate_potential))
 
 optimizer = optax.adam(learning_rate=0.01)
 
@@ -77,11 +67,7 @@ def calculate_potential(log_a, epsilon=0.5, sigma=1.0):
     history["a"].append(K.exp(log_a))
     history["energy"].append(energy)
 
-    # Check for NaN gradients using TensorCircuit's backend-agnostic approach
-    if K.sum(tc.num_to_tensor(np.isnan(K.numpy(grad)))) > 0:
-        print(f"Gradient became NaN at iteration {i+1}. Stopping optimization.")
-        print(f"Current energy: {energy}, Current log_a: {log_a}")
-        break
+    # (Removed previously added blanket NaN guard per reviewer request to keep example minimal.)
 
     updates, opt_state = optimizer.update(grad, opt_state)
     log_a = optax.apply_updates(log_a, updates)

tensorcircuit/backends/abstract_backend.py

@@ -629,15 +629,27 @@ def all(self: Any, a: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
 
     def meshgrid(self: Any, *args: Any, **kwargs: Any) -> Any:
         """
-        Return coordinate matrices from coordinate vectors.
+         Return coordinate matrices from coordinate vectors.
 
-        :param args: coordinate vectors
-        :type args: Any
+         :param args: coordinate vectors
+         :type args: Any
         :param kwargs: keyword arguments for meshgrid, typically includes 'indexing'
-            which can be 'ij' (matrix indexing) or 'xy' (Cartesian indexing)
-        :type kwargs: Any
-        :return: list of coordinate matrices
-        :rtype: Any
+           which can be 'ij' (matrix indexing) or 'xy' (Cartesian indexing).
+           - 'ij': matrix indexing, first dimension corresponds to rows (default)
+           - 'xy': Cartesian indexing, first dimension corresponds to columns
+           Example:
+              >>> x, y = backend.meshgrid([0, 1], [0, 2], indexing='xy')
+              Shapes:
+              - x.shape == (2, 2)  # rows correspond to y vector length
+              - y.shape == (2, 2)
+              Values:
+              x = [[0, 1],
+                  [0, 1]]
+              y = [[0, 0],
+                  [2, 2]]
+         :type kwargs: Any
+         :return: list of coordinate matrices
+         :rtype: Any
         """
         raise NotImplementedError(
             "Backend '{}' has not implemented `meshgrid`.".format(self.name)
@@ -797,6 +809,21 @@ def cast(self: Any, a: Tensor, dtype: str) -> Tensor:
             "Backend '{}' has not implemented `cast`.".format(self.name)
         )
 
+    def convert_to_tensor(self: Any, a: Tensor, dtype: Optional[str] = None) -> Tensor:
+        """
+        Convert input to tensor.
+
+        :param a: input data to be converted
+        :type a: Tensor
+        :param dtype: target dtype, optional
+        :type dtype: Optional[str]
+        :return: converted tensor
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `convert_to_tensor`.".format(self.name)
+        )
+
     def mod(self: Any, x: Tensor, y: Tensor) -> Tensor:
         """
         Compute y-mod of x (negative number behavior is not guaranteed to be consistent)

tensorcircuit/backends/pytorch_backend.py

@@ -253,8 +253,6 @@ def expm(self, a: Tensor) -> Tensor:
         # it doesn't support complex numbers which is more severe issue.
         # see https://github.com/pytorch/pytorch/issues/9983
 
-    # see https://github.com/pytorch/pytorch/issues/9983
-
     def sin(self, a: Tensor) -> Tensor:
         return torchlib.sin(a)
 

tensorcircuit/backends/tensorflow_backend.py

@@ -77,10 +77,8 @@ def _tensordot_tf(
 ) -> Tensor:
     # Use TensorFlow's dtype promotion rules by converting both to a common dtype
     if a.dtype != b.dtype:
-        # Find the result dtype by performing a dummy operation
-        common_dtype = (
-            tf.constant(0, dtype=a.dtype) + tf.constant(0, dtype=b.dtype)
-        ).dtype
+        # Find the result dtype using TensorFlow's type promotion rules
+        common_dtype = tf.experimental.numpy.result_type(a.dtype, b.dtype)
         a = tf.cast(a, common_dtype)
         b = tf.cast(b, common_dtype)
     return tf.tensordot(a, b, axes)