tensorcircuit-nightly 1.0.2.dev20250108__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/backends/abstract_backend.py

@@ -9,6 +9,7 @@ from functools import reduce, partial
 from operator import mul
 from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
 
+import math
 import numpy as np
 from ..utils import return_partial
 
@@ -46,17 +47,22 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `expm`.".format(self.name)
         )
 
-    def sqrtmh(self: Any, a: Tensor) -> Tensor:
+    def sqrtmh(self: Any, a: Tensor, psd: bool = False) -> Tensor:
         """
         Return the sqrtm of a Hermitian matrix ``a``.
 
         :param a: tensor in matrix form
         :type a: Tensor
+        :param psd: whether the input ``a`` is guaranteed as a positive semidefinite matrix,
+            defaults False
+        :type psd: bool
         :return: sqrtm of ``a``
         :rtype: Tensor
         """
         # maybe friendly for AD and also cosidering that several backend has no support for native sqrtm
         e, v = self.eigh(a)
+        if psd:
+            e = self.relu(e)
         e = self.sqrt(e)
         return v @ self.diagflat(e) @ self.adjoint(v)
 
@@ -400,6 +406,31 @@ class ExtendedBackend:
         a = self.reshape(a, [2 for _ in range(nleg)])
         return a
 
+    def reshaped(self: Any, a: Tensor, d: int) -> Tensor:
+        """
+        Reshape a tensor to the [d, d, ...] shape.
+
+        :param a: Input tensor
+        :type a: Tensor
+        :param d: edge length for each dimension
+        :type d: int
+        :return: the reshaped tensor
+        :rtype: Tensor
+        """
+        if not isinstance(d, int) or d <= 0:
+            raise ValueError("d must be a positive integer.")
+
+        size = self.sizen(a)
+        if size == 0:
+            return self.reshape(a, (0,))
+
+        nleg_float = math.log(size, d)
+        nleg = int(round(nleg_float))
+        if d**nleg != size:
+            raise ValueError(f"cannot reshape: size {size} is not a power of d={d}")
+
+        return self.reshape(a, (d,) * nleg)
+
     def reshapem(self: Any, a: Tensor) -> Tensor:
         """
         Reshape a tensor to the [l, l] shape.
@@ -576,6 +607,97 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `argmin`.".format(self.name)
         )
 
+    def argsort(self: Any, a: Tensor, axis: int = -1) -> Tensor:
+        """
+        return the indices that would sort an array.
+
+        :param a: the tensor to be sorted
+        :type a: Tensor
+        :param axis: the sorted axis, defaults to -1
+        :type axis: int
+        :return: the sorted indices
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "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``.
@@ -695,6 +817,9 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `is_tensor`.".format(self.name)
         )
 
+    def matvec(self: Any, A: Tensor, x: Tensor) -> Tensor:
+        return self.tensordot(A, x, axes=[[1], [0]])
+
     def cast(self: Any, a: Tensor, dtype: str) -> Tensor:
         """
         Cast the tensor dtype of a ``a``.
@@ -710,6 +835,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)
@@ -725,6 +865,82 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `mod`.".format(self.name)
         )
 
+    def floor_divide(self: Any, x: Tensor, y: Tensor) -> Tensor:
+        r"""
+        Compute the element-wise floor division of two tensors.
+
+        This operation returns a new tensor containing the result of
+        dividing `x` by `y` and rounding each element down towards
+        negative infinity. The semantics are equivalent to the Python
+        `//` operator:
+
+            result[i] = floor(x[i] / y[i])
+
+        Broadcasting is supported according to the backend's rules.
+
+        :param x: Dividend tensor.
+        :type x: Tensor
+        :param y: Divisor tensor, must be broadcastable with `x`.
+        :type y: Tensor
+        :return: A tensor with the broadcasted shape of `x` and `y`,
+                 where each element is the floored result of the division.
+        :rtype: Tensor
+
+        :raises NotImplementedError: If the backend does not provide an
+                                     implementation for `floor_divide`.
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `floor_divide`.".format(self.name)
+        )
+
+    def floor(self: Any, a: Tensor) -> Tensor:
+        """
+        Compute the element-wise floor of the input tensor.
+
+        This operation returns a new tensor with the largest integers
+        less than or equal to each element of the input tensor,
+        i.e. it rounds each value down towards negative infinity.
+
+        :param a: Input tensor containing numeric values.
+        :type a: Tensor
+        :return: A tensor with the same shape as `a`, where each element
+                 is the floored value of the corresponding element in `a`.
+        :rtype: Tensor
+
+        :raises NotImplementedError: If the backend does not provide an
+                                     implementation for `floor`.
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `floor`.".format(self.name)
+        )
+
+    def clip(self: Any, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        """
+        Clip (limit) the values of a tensor element-wise to the range [a_min, a_max].
+
+        Each element in the input tensor `a` is compared against the corresponding
+        bounds `a_min` and `a_max`. If a value in `a` is less than `a_min`, it is set
+        to `a_min`; if greater than `a_max`, it is set to `a_max`. Otherwise, the
+        value is left unchanged. The result preserves the dtype and device of the input.
+
+        :param a: Input tensor containing values to be clipped.
+        :type a: Tensor
+        :param a_min: Lower bound (minimum value) for clipping. Can be a scalar tensor
+                      or broadcastable to the shape of `a`.
+        :type a_min: Tensor
+        :param a_max: Upper bound (maximum value) for clipping. Can be a scalar tensor
+                      or broadcastable to the shape of `a`.
+        :type a_max: Tensor
+        :return: A tensor with the same shape as `a`, where all values are clipped
+                 to lie within the interval [a_min, a_max].
+        :rtype: Tensor
+
+        :raises NotImplementedError: If the backend does not implement `clip`.
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `clip`.".format(self.name)
+        )
+
     def reverse(self: Any, a: Tensor) -> Tensor:
         """
         return ``a[::-1]``, only 1D tensor is guaranteed for consistent behavior
@@ -800,6 +1016,21 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `solve`.".format(self.name)
         )
 
+    def special_jv(self: Any, v: int, z: Tensor, M: int) -> Tensor:
+        """
+        Special function: Bessel function of the first kind.
+
+        :param v: The order of the Bessel function.
+        :type v: int
+        :param z: The argument of the Bessel function.
+        :type z: Tensor
+        :return: The value of the Bessel function [J_0, ...J_{v-1}(z)].
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `special_jv`.".format(self.name)
+        )
+
     def searchsorted(self: Any, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         """
         Find indices where elements should be inserted to maintain order.
@@ -808,8 +1039,8 @@ class ExtendedBackend:
         :type a: Tensor
         :param v: value to inserted
         :type v: Tensor
-        :param side:  If ‘left’, the index of the first suitable location found is given.
-            If ‘right’, return the last such index.
+        :param side:  If `left`, the index of the first suitable location found is given.
+            If `right`, return the last such index.
             If there is no suitable index, return either 0 or N (where N is the length of a),
             defaults to "left"
         :type side: str, optional
@@ -1243,6 +1474,26 @@ class ExtendedBackend:
             "Backend '{}' has not implemented `sparse_dense_matmul`.".format(self.name)
         )
 
+    def sparse_csr_from_coo(self: Any, coo: Tensor, strict: bool = False) -> Tensor:
+        """
+        transform a coo matrix to a csr matrix
+
+        :param coo: a coo matrix
+        :type coo: Tensor
+        :param strict: whether to enforce the transform, defaults to False,
+            corresponding to return the coo matrix if there is no implementation for specific backend.
+        :type strict: bool, optional
+        :return: a csr matrix
+        :rtype: Tensor
+        """
+        if strict:
+            raise NotImplementedError(
+                "Backend '{}' has not implemented `sparse_csr_from_coo`.".format(
+                    self.name
+                )
+            )
+        return coo
+
     def to_dense(self: Any, sp_a: Tensor) -> Tensor:
         """
         Convert a sparse matrix to dense tensor.
@@ -1346,6 +1597,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:
@@ -1381,10 +1654,49 @@ class ExtendedBackend:
         :rtype: Tensor
         """
         carry = init
-        for x in xs:
-            carry = f(carry, x)
+        # Check if `xs` is a PyTree (tuple or list) of arrays.
+        if isinstance(xs, (tuple, list)):
+            for x_slice_tuple in zip(*xs):
+                # x_slice_tuple will be (k_elems[i], j_elems[i]) at each step.
+                carry = f(carry, x_slice_tuple)
+        else:
+            # If xs is a single array, iterate normally.
+            for x in xs:
+                carry = f(carry, x)
+
         return carry
 
+    def jaxy_scan(
+        self: Any, f: Callable[[Tensor, Tensor], Tensor], init: Tensor, xs: Tensor
+    ) -> Tensor:
+        """
+        This API follows jax scan style. TF use plain for loop
+
+        :param f: _description_
+        :type f: Callable[[Tensor, Tensor], Tensor]
+        :param init: _description_
+        :type init: Tensor
+        :param xs: _description_
+        :type xs: Tensor
+        :raises ValueError: _description_
+        :return: _description_
+        :rtype: Tensor
+        """
+        if xs is None:
+            raise ValueError("Either xs or length must be provided.")
+        if xs is not None:
+            length = len(xs)
+        carry, outputs_to_stack = init, []
+        for i in range(length):
+            if isinstance(xs, (tuple, list)):
+                x = [ele[i] for ele in xs]
+            else:
+                x = xs[i]
+            new_carry, y = f(carry, x)
+            carry = new_carry
+            outputs_to_stack.append(y)
+        return carry, self.stack(outputs_to_stack)
+
     def stop_gradient(self: Any, a: Tensor) -> Tensor:
         """
         Stop backpropagation from ``a``.
@@ -1548,7 +1860,7 @@ class ExtendedBackend:
                                 if i == argnum
                                 else self.reshape(
                                     self.zeros(
-                                        [self.sizen(arg), self.sizen(arg)],
+                                        [self.sizen(args[argnum]), self.sizen(arg)],
                                         dtype=arg.dtype,
                                     ),
                                     [-1] + list(self.shape_tuple(arg)),
@@ -1636,6 +1948,7 @@ class ExtendedBackend:
                     ),
                     jj,
                 )
+                jj = [jji for ind, jji in enumerate(jj) if ind in argnums]
                 if len(jj) == 1:
                     jj = jj[0]
                 jjs.append(jj)
@@ -1662,7 +1975,7 @@ class ExtendedBackend:
         f: Callable[..., Any],
         static_argnums: Optional[Union[int, Sequence[int]]] = None,
         jit_compile: Optional[bool] = None,
-        **kws: Any
+        **kws: Any,
     ) -> Callable[..., Any]:
         """
         Return the jitted version of function ``f``.

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
 
 
@@ -170,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
 )
@@ -219,6 +245,11 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
 
         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
@@ -243,8 +274,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:
@@ -342,6 +375,15 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return jnp.mod(x, y)
 
+    def floor(self, a: Tensor) -> Tensor:
+        return jnp.floor(a)
+
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return jnp.floor_divide(x, y)
+
+    def clip(self, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        return jnp.clip(a, a_min, a_max)
+
     def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
         return jnp.right_shift(x, y)
 
@@ -387,6 +429,12 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return jnp.argmin(a, axis=axis)
 
+    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]:
@@ -407,6 +455,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
@@ -418,6 +469,11 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:  # type: ignore
         return jsp.linalg.solve(A, b, assume_a=assume_a)
 
+    def special_jv(self, v: int, z: Tensor, M: int) -> Tensor:
+        from .jax_ops import bessel_jv_jax_rescaled
+
+        return bessel_jv_jax_rescaled(v, z, M)
+
     def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         if not self.is_tensor(a):
             a = self.convert_to_tensor(a)
@@ -442,7 +498,14 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     def to_dlpack(self, a: Tensor) -> Any:
         import jax.dlpack
 
-        return jax.dlpack.to_dlpack(a)
+        try:
+            return jax.dlpack.to_dlpack(a)  # type: ignore
+        except AttributeError:  # jax >v0.7
+            # jax.dlpack.to_dlpack was deprecated in JAX v0.6.0 and removed in JAX v0.7.0.
+            # Please use the newer DLPack API based on __dlpack__ and __dlpack_device__ instead.
+            # Typically, you can pass a JAX array directly to the `from_dlpack` function of
+            # another framework without using `to_dlpack`.
+            return a.__dlpack__()
 
     def set_random_state(
         self, seed: Optional[Union[int, PRNGKeyArray]] = None, get_only: bool = False
@@ -608,7 +671,14 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
         carry, _ = libjax.lax.scan(f_jax, init, xs)
         return carry
 
+    def jaxy_scan(
+        self, f: Callable[[Tensor, Tensor], Tensor], init: Tensor, xs: Tensor
+    ) -> Tensor:
+        return libjax.lax.scan(f, init, xs)
+
     def scatter(self, operand: Tensor, indices: Tensor, updates: Tensor) -> Tensor:
+        # updates = jnp.reshape(updates, indices.shape)
+        # return operand.at[indices].set(updates)
         rank = len(operand.shape)
         dnums = libjax.lax.ScatterDimensionNumbers(
             update_window_dims=(),
@@ -630,11 +700,20 @@ class JaxBackend(jax_backend.JaxBackend, ExtendedBackend):  # type: ignore
     ) -> Tensor:
         return sp_a @ b
 
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        try:
+            return sparse.BCSR.from_bcoo(coo)  # type: ignore
+        except AttributeError as e:
+            if not strict:
+                return coo
+            else:
+                raise e
+
     def to_dense(self, sp_a: Tensor) -> Tensor:
         return sp_a.todense()
 
     def is_sparse(self, a: Tensor) -> bool:
-        return isinstance(a, sparse.BCOO)  # type: ignore
+        return isinstance(a, sparse.JAXSparse)  # type: ignore
 
     def device(self, a: Tensor) -> str:
         (dev,) = a.devices()
@@ -781,4 +860,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)