tensorcircuit-nightly 1.3.0.dev20250815__py3-none-any.whl → 1.3.0.dev20250816__py3-none-any.whl

tensorcircuit/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.3.0.dev20250815"
+__version__ = "1.3.0.dev20250816"
 __author__ = "TensorCircuit Authors"
 __creator__ = "refraction-ray"
 

tensorcircuit/backends/abstract_backend.py

@@ -596,6 +596,82 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `argsort`.".format(self.name)
         )
 
+    def sort(self: Any, a: Tensor, axis: int = -1) -> Tensor:
+        """
+        Sort a tensor along the given axis.
+
+        :param a: [description]
+        :type a: Tensor
+        :param axis: [description], defaults to -1
+        :type axis: int, optional
+        :return: [description]
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `sort`.".format(self.name)
+        )
+
+    def all(self: Any, a: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
+        """
+        Test whether all array elements along a given axis evaluate to True.
+
+        :param a: Input tensor
+        :type a: Tensor
+        :param axis: Axis or axes along which a logical AND reduction is performed,
+            defaults to None
+        :type axis: Optional[Sequence[int]], optional
+        :return: A new boolean or tensor resulting from the AND reduction
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `all`.".format(self.name)
+        )
+
+    def meshgrid(self: Any, *args: Any, **kwargs: Any) -> Any:
+        """
+        Return coordinate matrices from coordinate vectors.
+
+        :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).
+           - '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)
+        )
+
+    def expand_dims(self: Any, a: Tensor, axis: int) -> Tensor:
+        """
+        Expand the shape of a tensor.
+        Insert a new axis that will appear at the `axis` position in the expanded
+        tensor shape.
+
+        :param a: Input tensor
+        :type a: Tensor
+        :param axis: Position in the expanded axes where the new axis is placed
+        :type axis: int
+        :return: Output tensor with the number of dimensions increased by one.
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `expand_dims`.".format(self.name)
+        )
+
     def unique_with_counts(self: Any, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         """
         Find the unique elements and their corresponding counts of the given tensor ``a``.
@@ -733,6 +809,21 @@ class ExtendedBackend:
             "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)
@@ -1404,6 +1495,28 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `cond`.".format(self.name)
         )
 
+    def where(
+        self: Any,
+        condition: Tensor,
+        x: Optional[Tensor] = None,
+        y: Optional[Tensor] = None,
+    ) -> Tensor:
+        """
+        Return a tensor of elements selected from either x or y, depending on condition.
+
+        :param condition: Where True, yield x, otherwise yield y.
+        :type condition: Tensor (bool)
+        :param x: Values from which to choose when condition is True.
+        :type x: Tensor
+        :param y: Values from which to choose when condition is False.
+        :type y: Tensor
+        :return: A tensor with elements from x where condition is True, and y otherwise.
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `where`.".format(self.name)
+        )
+
     def switch(
         self: Any, index: Tensor, branches: Sequence[Callable[[], Tensor]]
     ) -> Tensor:

tensorcircuit/backends/cupy_backend.py

@@ -56,10 +56,12 @@ class CuPyBackend(tnbackend, ExtendedBackend):  # type: ignore
         cpx = cupyx
         self.name = "cupy"
 
-    def convert_to_tensor(self, a: Tensor) -> Tensor:
+    def convert_to_tensor(self, a: Tensor, dtype: Optional[str] = None) -> Tensor:
         if not isinstance(a, cp.ndarray) and not cp.isscalar(a):
             a = cp.array(a)
         a = cp.asarray(a)
+        if dtype is not None:
+            a = self.cast(a, dtype)
         return a
 
     def sum(

tensorcircuit/backends/jax_backend.py

@@ -50,12 +50,17 @@ class optax_optimizer:
         return params
 
 
-def _convert_to_tensor_jax(self: Any, tensor: Tensor) -> Tensor:
+def _convert_to_tensor_jax(
+    self: Any, tensor: Tensor, dtype: Optional[str] = None
+) -> Tensor:
     if not isinstance(tensor, (np.ndarray, jnp.ndarray)) and not jnp.isscalar(tensor):
         raise TypeError(
             ("Expected a `jnp.array`, `np.array` or scalar. " f"Got {type(tensor)}")
         )
     result = jnp.asarray(tensor)
+    if dtype is not None:
+        # Use the backend's cast method to handle dtype conversion
+        result = self.cast(result, dtype)
     return result
 
 
@@ -243,8 +248,10 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def copy(self, tensor: Tensor) -> Tensor:
         return jnp.array(tensor, copy=True)
 
-    def convert_to_tensor(self, tensor: Tensor) -> Tensor:
+    def convert_to_tensor(self, tensor: Tensor, dtype: Optional[str] = None) -> Tensor:
         result = jnp.asarray(tensor)
+        if dtype is not None:
+            result = self.cast(result, dtype)
         return result
 
     def abs(self, a: Tensor) -> Tensor:
@@ -390,6 +397,9 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
         return jnp.argsort(a, axis=axis)
 
+    def sort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return jnp.sort(a, axis=axis)
+
     def unique_with_counts(  # type: ignore
         self, a: Tensor, *, size: Optional[int] = None, fill_value: Optional[int] = None
     ) -> Tuple[Tensor, Tensor]:
@@ -410,6 +420,9 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def cumsum(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
         return jnp.cumsum(a, axis)
 
+    def all(self, a: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
+        return jnp.all(a, axis=axis)
+
     def is_tensor(self, a: Any) -> bool:
         if not isinstance(a, jnp.ndarray):
             return False
@@ -812,4 +825,23 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
 
     vvag = vectorized_value_and_grad
 
+    def meshgrid(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Backend-agnostic meshgrid function.
+        """
+        return jnp.meshgrid(*args, **kwargs)
+
     optimizer = optax_optimizer
+
+    def expand_dims(self, a: Tensor, axis: int) -> Tensor:
+        return jnp.expand_dims(a, axis)
+
+    def where(
+        self,
+        condition: Tensor,
+        x: Optional[Tensor] = None,
+        y: Optional[Tensor] = None,
+    ) -> Tensor:
+        if x is None and y is None:
+            return jnp.where(condition)
+        return jnp.where(condition, x, y)

tensorcircuit/backends/numpy_backend.py

@@ -35,10 +35,14 @@ def _sum_numpy(
     # see https://github.com/google/TensorNetwork/issues/952
 
 
-def _convert_to_tensor_numpy(self: Any, a: Tensor) -> Tensor:
+def _convert_to_tensor_numpy(
+    self: Any, a: Tensor, dtype: Optional[str] = None
+) -> Tensor:
     if not isinstance(a, np.ndarray) and not np.isscalar(a):
         a = np.array(a)
     a = np.asarray(a)
+    if dtype is not None:
+        a = a.astype(getattr(np, dtype))
     return a
 
 
@@ -132,6 +136,9 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
     def kron(self, a: Tensor, b: Tensor) -> Tensor:
         return np.kron(a, b)
 
+    def meshgrid(self, *args: Any, **kwargs: Any) -> Any:
+        return np.meshgrid(*args, **kwargs)
+
     def dtype(self, a: Tensor) -> str:
         return a.dtype.__str__()  # type: ignore
 
@@ -151,6 +158,9 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
             dtype = getattr(np, dtype)
         return np.array(1j, dtype=dtype)
 
+    def expand_dims(self, a: Tensor, axis: int) -> Tensor:
+        return np.expand_dims(a, axis)
+
     def stack(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
         return np.stack(a, axis=axis)
 
@@ -173,6 +183,9 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
     ) -> Tensor:
         return np.std(a, axis=axis, keepdims=keepdims)
 
+    def all(self, a: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
+        return np.all(a, axis=axis)
+
     def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         return np.unique(a, return_counts=True)  # type: ignore
 
@@ -188,6 +201,9 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return np.argmin(a, axis=axis)
 
+    def sort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return np.sort(a, axis=axis)
+
     def sigmoid(self, a: Tensor) -> Tensor:
         return expit(a)
 
@@ -345,6 +361,17 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
     def is_sparse(self, a: Tensor) -> bool:
         return issparse(a)  # type: ignore
 
+    def where(
+        self,
+        condition: Tensor,
+        x: Optional[Tensor] = None,
+        y: Optional[Tensor] = None,
+    ) -> Tensor:
+        if x is None and y is None:
+            return np.where(condition)
+        assert x is not None and y is not None
+        return np.where(condition, x, y)
+
     def cond(
         self,
         pred: bool,

tensorcircuit/backends/pytorch_backend.py

@@ -238,6 +238,15 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
     def copy(self, a: Tensor) -> Tensor:
         return a.clone()
 
+    def convert_to_tensor(self, tensor: Tensor, dtype: Optional[str] = None) -> Tensor:
+        if self.is_tensor(tensor):
+            result = tensor
+        else:
+            result = torchlib.tensor(tensor)
+        if dtype is not None:
+            result = self.cast(result, dtype)
+        return result
+
     def expm(self, a: Tensor) -> Tensor:
         raise NotImplementedError("pytorch backend doesn't support expm")
         # in 2020, torch has no expm, hmmm. but that's ok,
@@ -369,6 +378,17 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return torchlib.argmin(a, dim=axis)
 
+    def sort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return torchlib.sort(a, dim=axis).values
+
+    def all(self, tensor: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
+        """
+        Corresponds to torch.all.
+        """
+        if axis is None:
+            return torchlib.all(tensor)
+        return torchlib.all(tensor, dim=axis)
+
     def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         return torchlib.unique(a, return_counts=True)  # type: ignore
 
@@ -425,6 +445,16 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
             v = self.convert_to_tensor(v)
         return torchlib.searchsorted(a, v, side=side)
 
+    def where(
+        self,
+        condition: Tensor,
+        x: Optional[Tensor] = None,
+        y: Optional[Tensor] = None,
+    ) -> Tensor:
+        if x is None and y is None:
+            return torchlib.where(condition)
+        return torchlib.where(condition, x, y)
+
     def reverse(self, a: Tensor) -> Tensor:
         return torchlib.flip(a, dims=(-1,))
 
@@ -706,6 +736,12 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
 
         return wrapper
 
+    def expand_dims(self, a: Tensor, axis: int) -> Tensor:
+        return torchlib.unsqueeze(a, dim=axis)
+
     vvag = vectorized_value_and_grad
 
+    def meshgrid(self, *args: Any, **kws: Any) -> Tensor:
+        return torchlib.meshgrid(*args, **kws)
+
     optimizer = torch_optimizer

tensorcircuit/backends/tensorflow_backend.py

@@ -75,6 +75,12 @@ class keras_optimizer:
 def _tensordot_tf(
     self: Any, a: Tensor, b: Tensor, axes: Union[int, Sequence[Sequence[int]]]
 ) -> Tensor:
+    # Use TensorFlow's dtype promotion rules by converting both to a common dtype
+    if a.dtype != b.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)
 
 
@@ -441,6 +447,12 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
     def copy(self, a: Tensor) -> Tensor:
         return tf.identity(a)
 
+    def convert_to_tensor(self, tensor: Tensor, dtype: Optional[str] = None) -> Tensor:
+        result = tf.convert_to_tensor(tensor)
+        if dtype is not None:
+            result = self.cast(result, dtype)
+        return result
+
     def expm(self, a: Tensor) -> Tensor:
         return tf.linalg.expm(a)
 
@@ -524,6 +536,20 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
     def max(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
         return tf.reduce_max(a, axis=axis)
 
+    def all(self, a: Tensor, axis: Optional[Sequence[int]] = None) -> Tensor:
+        return tf.reduce_all(tf.cast(a, tf.bool), axis=axis)
+
+    def where(
+        self,
+        condition: Tensor,
+        x: Optional[Tensor] = None,
+        y: Optional[Tensor] = None,
+    ) -> Tensor:
+        if x is None and y is None:
+            # Return a tuple of tensors to be consistent with other backends
+            return tuple(tf.unstack(tf.where(condition), axis=1))
+        return tf.where(condition, x, y)
+
     def argmax(self, a: Tensor, axis: int = 0) -> Tensor:
         return tf.math.argmax(a, axis=axis)
 
@@ -533,6 +559,9 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
     def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
         return tf.argsort(a, axis=axis)
 
+    def sort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return tf.sort(a, axis=axis)
+
     def shape_tuple(self, a: Tensor) -> Tuple[int, ...]:
         return tuple(a.shape)
 
@@ -1061,4 +1090,13 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
 
     vvag = vectorized_value_and_grad
 
+    def meshgrid(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Backend-agnostic meshgrid function.
+        """
+        return tf.meshgrid(*args, **kwargs)
+
     optimizer = keras_optimizer
+
+    def expand_dims(self, a: Tensor, axis: int) -> Tensor:
+        return tf.expand_dims(a, axis)

tensorcircuit/templates/hamiltonians.py

@@ -17,13 +17,14 @@ def _create_empty_sparse_matrix(shape: Tuple[int, int]) -> Any:
 def heisenberg_hamiltonian(
     lattice: AbstractLattice,
     j_coupling: Union[float, List[float], Tuple[float, ...]] = 1.0,
+    interaction_scope: str = "neighbors",
 ) -> Any:
     """
     Generates the sparse matrix of the Heisenberg Hamiltonian for a given lattice.
 
     The Heisenberg Hamiltonian is defined as:
-    H = J * Σ_{<i,j>} (X_i X_j + Y_i Y_j + Z_i Z_j)
-    where the sum is over all unique nearest-neighbor pairs <i,j>.
+    H = J * Σ_{i,j} (X_i X_j + Y_i Y_j + Z_i Z_j)
+    where the sum is over a specified set of interacting pairs {i,j}.
 
     :param lattice: An instance of a class derived from AbstractLattice,
         which provides the geometric information of the system.
@@ -32,11 +33,23 @@ def heisenberg_hamiltonian(
         isotropic model (Jx=Jy=Jz) or a list/tuple of 3 floats for an
         anisotropic model (Jx, Jy, Jz). Defaults to 1.0.
     :type j_coupling: Union[float, List[float], Tuple[float, ...]], optional
+    :param interaction_scope: Defines the range of interactions.
+        - "neighbors": Includes only nearest-neighbor pairs (default).
+        - "all": Includes all unique pairs of sites.
+    :type interaction_scope: str, optional
     :return: The Hamiltonian as a backend-agnostic sparse matrix.
     :rtype: Any
     """
     num_sites = lattice.num_sites
-    neighbor_pairs = lattice.get_neighbor_pairs(k=1, unique=True)
+    if interaction_scope == "neighbors":
+        neighbor_pairs = lattice.get_neighbor_pairs(k=1, unique=True)
+    elif interaction_scope == "all":
+        neighbor_pairs = lattice.get_all_pairs()
+    else:
+        raise ValueError(
+            f"Invalid interaction_scope: '{interaction_scope}'. "
+            "Must be 'neighbors' or 'all'."
+        )
 
     if isinstance(j_coupling, (float, int)):
         js = [float(j_coupling)] * 3