tensorcircuit-nightly 1.3.0.dev20250728__py3-none-any.whl → 1.4.0.dev20251103__py3-none-any.whl

tensorcircuit/backends/jax_ops.py

@@ -3,8 +3,11 @@ Customized ops for ML framework
 """
 
 # pylint: disable=invalid-name
+# pylint: disable=unused-variable
+
 
 from typing import Any, Tuple, Sequence
+from functools import partial
 
 import jax
 import jax.numpy as jnp
@@ -174,3 +177,108 @@ def jaxeigh_bwd(r: Array, tangents: Array) -> Array:
 
 adaware_eigh.defvjp(jaxeigh_fwd, jaxeigh_bwd)
 adaware_eigh_jit = jax.jit(adaware_eigh)
+
+
+@partial(jax.jit, static_argnums=[0, 2])
+def bessel_jv_jax_rescaled(k: int, x: jnp.ndarray, M: int) -> jnp.ndarray:
+    """
+    Computes Bessel function Jv using Miller's algorithm with dynamic rescaling,
+    implemented in JAX.
+    """
+    if M <= k:
+        raise ValueError(
+            f"Recurrence length M ({M}) must be greater than the required order k ({k})."
+        )
+
+    # Use vmap to handle array inputs for x efficiently.
+    # We map _bessel_jv_scalar_rescaled over the last dimension of x.
+    return _bessel_jv_scalar_rescaled(k, M, x)
+
+
+def _bessel_jv_scalar_rescaled(k: int, M: int, x_val: jnp.ndarray) -> jnp.ndarray:
+    """
+    JAX implementation for a scalar input x_val.
+    This function will be vmapped for array inputs.
+    """
+    rescale_threshold = 1e250
+
+    # Define the body of the recurrence loop
+    def recurrence_body(i, state):  # type: ignore
+        # M - i is the current 'm' value in the original loop.
+        # Loop from M down to 1. jax.lax.fori_loop goes from lower to upper-1.
+        # So for m from M down to 1, we map i from 0 to M-1.
+        # Current m_val = M - i
+        # The loop range for m in numpy was `range(M, 0, -1)`, which means m goes from M, M-1, ..., 1.
+        # For lax.fori_loop (start, stop, body_fn, init_val), start is inclusive, stop is exclusive.
+        # So to iterate M times for m from M down to 1, we do i from 0 to M-1.
+        # m_val = M - i means that for i=0, m_val=M; for i=M-1, m_val=1.
+        m_val = M - i
+        f_m, f_m_p1, f_vals = state
+
+        # If x_val is near zero, this division could be an issue,
+        # but the outer lax.cond handles the x_val near zero case before this loop runs.
+        f_m_m1 = (2.0 * m_val / x_val) * f_m - f_m_p1
+
+        # --- Rescaling Step ---
+        # jax.lax.cond requires all branches to return the exact same type and shape.
+        def rescale_branch(vals):  # type: ignore
+            f_m_val, f_m_p1_val, f_vals_arr = vals
+            scale_factor = f_m_m1
+            # Return new f_m, f_m_p1, updated f_vals_arr, and the new f_m_m1 value (which is 1.0)
+            return (
+                f_m_val / scale_factor,
+                f_m_p1_val / scale_factor,
+                f_vals_arr / scale_factor,
+                1.0,
+            )
+
+        def no_rescale_branch(vals):  # type: ignore
+            f_m_val, f_m_p1_val, f_vals_arr = (
+                vals  # Unpack to keep signatures consistent
+            )
+            # Return original f_m, f_m_p1, original f_vals_arr, and the computed f_m_m1
+            return (f_m_val, f_m_p1_val, f_vals_arr, f_m_m1)
+
+        f_m_rescaled, f_m_p1_rescaled, f_vals_rescaled, f_m_m1_effective = jax.lax.cond(
+            jnp.abs(f_m_m1) > rescale_threshold,
+            rescale_branch,
+            no_rescale_branch,
+            (f_m, f_m_p1, f_vals),  # Arguments passed to branches
+        )
+
+        # Update f_vals at index m_val - 1. JAX uses .at[idx].set(val) for non-in-place updates.
+        f_vals_updated = f_vals_rescaled.at[m_val - 1].set(f_m_m1_effective)
+
+        # Return new state for the next iteration: (new f_m, new f_m_p1, updated f_vals)
+        return (f_m_m1_effective, f_m_rescaled, f_vals_updated)
+
+    # Initial state for the recurrence loop
+    f_m_p1_init = 0.0
+    f_m_init = 1e-30  # Start with a very small number
+    f_vals_init = jnp.zeros(M + 1).at[M].set(f_m_init)
+
+    # Use jax.lax.fori_loop for the backward recurrence
+    # Loop from i = 0 to M-1 (total M iterations)
+    # The 'body' function gets current 'i' and 'state', returns 'new_state'.
+    # We don't need the final f_m_p1, only f_m and f_vals.
+    final_f_m, _, f_vals = jax.lax.fori_loop(
+        0, M, recurrence_body, (f_m_init, f_m_p1_init, f_vals_init)
+    )
+
+    # Normalization using Neumann's sum rule
+    even_sum = jnp.sum(f_vals[2::2])
+    norm_const = f_vals[0] + 2.0 * even_sum
+
+    # Handle division by near-zero normalization constant
+    norm_const_safe = jnp.where(jnp.abs(norm_const) < 1e-12, 1e-12, norm_const)
+
+    # Conditional logic for x_val close to zero
+    def x_is_zero_case() -> jnp.ndarray:
+        # For x=0, J_0(0)=1, J_k(0)=0 for k>0
+        return jnp.zeros(k).at[0].set(1.0)
+
+    def x_is_not_zero_case() -> jnp.ndarray:
+        return f_vals[:k] / norm_const_safe  # type: ignore
+
+    # Use lax.cond to select between the two cases based on x_val
+    return jax.lax.cond(jnp.abs(x_val) < 1e-12, x_is_zero_case, x_is_not_zero_case)  # type: ignore

tensorcircuit/backends/numpy_backend.py

@@ -17,7 +17,7 @@ except ImportError:  # np2.0 compatibility
 
 import tensornetwork
 from scipy.linalg import expm, solve, schur
-from scipy.special import softmax, expit
+from scipy.special import softmax, expit, jv
 from scipy.sparse import coo_matrix, issparse
 from tensornetwork.backends.numpy import numpy_backend
 from .abstract_backend import ExtendedBackend
@@ -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)
 
@@ -200,6 +216,7 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
         return softmax(a, axis=axis)
 
     def onehot(self, a: Tensor, num: int) -> Tensor:
+        a = np.asarray(a)
         res = np.eye(num)[a.reshape([-1])]
         return res.reshape(list(a.shape) + [num])
         # https://stackoverflow.com/questions/38592324/one-hot-encoding-using-numpy
@@ -233,6 +250,15 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return np.mod(x, y)
 
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return np.floor_divide(x, y)
+
+    def floor(self, a: Tensor) -> Tensor:
+        return np.floor(a)
+
+    def clip(self, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        return np.clip(a, a_min, a_max)
+
     def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
         return np.right_shift(x, y)
 
@@ -244,9 +270,15 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
         # https://stackoverflow.com/questions/44672029/difference-between-numpy-linalg-solve-and-numpy-linalg-lu-solve/44710451
         return solve(A, b, assume_a=assume_a)
 
+    def special_jv(self, v: int, z: Tensor, M: int) -> Tensor:
+        return jv(np.arange(v), z)
+
     def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
         return np.searchsorted(a, v, side=side)  # type: ignore
 
+    def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return np.argsort(a, axis=axis)
+
     def set_random_state(
         self, seed: Optional[int] = None, get_only: bool = False
     ) -> Any:
@@ -329,12 +361,26 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
     ) -> Tensor:
         return sp_a @ b
 
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        return coo.tocsr()
+
     def to_dense(self, sp_a: Tensor) -> Tensor:
         return sp_a.todense()
 
     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,
@@ -390,7 +436,7 @@ class NumpyBackend(numpy_backend.NumPyBackend, ExtendedBackend):  # type: ignore
         f: Callable[..., Any],
         static_argnums: Optional[Union[int, Sequence[int]]] = None,
         jit_compile: Optional[bool] = None,
-        **kws: Any
+        **kws: Any,
     ) -> Callable[..., Any]:
         logger.info("numpy backend has no jit interface, just do nothing")
         return f

tensorcircuit/backends/pytorch_backend.py

@@ -9,6 +9,7 @@ from typing import Any, Callable, Optional, Sequence, Tuple, Union
 from operator import mul
 from functools import reduce, partial
 
+from scipy.sparse import coo_matrix
 import tensornetwork
 from tensornetwork.backends.pytorch import pytorch_backend
 from .abstract_backend import ExtendedBackend
@@ -23,7 +24,6 @@ logger = logging.getLogger(__name__)
 
 # TODO(@refraction-ray): lack stateful random methods implementation for now
 # TODO(@refraction-ray): lack scatter impl for now
-# TODO(@refraction-ray): lack sparse relevant methods for now
 # To be added once pytorch backend is ready
 
 
@@ -229,6 +229,9 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
         r = torchlib.ones(shape)
         return self.cast(r, dtype)
 
+    def exp(self, tensor: Tensor) -> Tensor:
+        return torchlib.exp(tensor)
+
     def zeros(self, shape: Tuple[int, ...], dtype: Optional[str] = None) -> Tensor:
         if dtype is None:
             dtype = dtypestr
@@ -238,8 +241,18 @@ 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")
+        return torchlib.linalg.matrix_exp(a)
+        # raise NotImplementedError("pytorch backend doesn't support expm")
         # in 2020, torch has no expm, hmmm. but that's ok,
         # it doesn't support complex numbers which is more severe issue.
         # see https://github.com/pytorch/pytorch/issues/9983
@@ -293,6 +306,9 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
         return torchlib.kron(a, b)
 
     def numpy(self, a: Tensor) -> Tensor:
+        if self.is_sparse(a):
+            a = a.coalesce()
+            return coo_matrix((a.values().numpy(), a.indices().numpy()), shape=a.shape)
         a = a.cpu()
         if a.is_conj():
             return a.resolve_conj().numpy()
@@ -369,6 +385,20 @@ 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 argsort(self, a: Tensor, axis: int = -1) -> Tensor:
+        return torchlib.argsort(a, dim=axis)
+
+    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
 
@@ -382,6 +412,7 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
         return torchlib.nn.Softmax(a, dim=axis)
 
     def onehot(self, a: Tensor, num: int) -> Tensor:
+        a = a.long()
         return torchlib.nn.functional.one_hot(a, num)
 
     def cumsum(self, a: Tensor, axis: Optional[int] = None) -> Tensor:
@@ -409,6 +440,15 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
     def mod(self, x: Tensor, y: Tensor) -> Tensor:
         return torchlib.fmod(x, y)
 
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return torchlib.floor_divide(x, y)
+
+    def floor(self, a: Tensor) -> Tensor:
+        return torchlib.floor(a)
+
+    def clip(self, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        return torchlib.clamp(a, a_min, a_max)
+
     def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
         return torchlib.bitwise_right_shift(x, y)
 
@@ -425,9 +465,52 @@ 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,))
 
+    def coo_sparse_matrix(
+        self, indices: Tensor, values: Tensor, shape: Tensor
+    ) -> Tensor:
+        # Convert COO format to PyTorch sparse tensor
+        indices = self.convert_to_tensor(indices)
+        return torchlib.sparse_coo_tensor(self.transpose(indices), values, shape)
+
+    def sparse_dense_matmul(
+        self,
+        sp_a: Tensor,
+        b: Tensor,
+    ) -> Tensor:
+        # Matrix multiplication between sparse and dense tensor
+        return torchlib.sparse.mm(sp_a, b)
+
+    def sparse_csr_from_coo(self, coo: Tensor, strict: bool = False) -> Tensor:
+        try:
+            # Convert COO to CSR format if supported
+            return coo.to_sparse_csr()
+        except AttributeError as e:
+            if not strict:
+                return coo
+            else:
+                raise e
+
+    def to_dense(self, sp_a: Tensor) -> Tensor:
+        # Convert sparse tensor to dense
+        return sp_a.to_dense()
+
+    def is_sparse(self, a: Tensor) -> bool:
+        # Check if tensor is sparse
+        return a.is_sparse or a.is_sparse_csr  # type: ignore
+
     def tree_map(self, f: Callable[..., Any], *pytrees: Any) -> Any:
         # torch native tree_map not support multiple pytree args
         # return torchlib.utils._pytree.tree_map(f, *pytrees)
@@ -661,7 +744,7 @@ class PyTorchBackend(pytorch_backend.PyTorchBackend, ExtendedBackend):  # type:
         f: Callable[..., Any],
         static_argnums: Optional[Union[int, Sequence[int]]] = None,
         jit_compile: Optional[bool] = None,
-        **kws: Any
+        **kws: Any,
     ) -> Any:
         if jit_compile is True:
             # experimental feature reusing the jit_compile flag for tf
@@ -706,6 +789,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)
 
 
@@ -360,6 +366,38 @@ tensornetwork.backends.tensorflow.tensorflow_backend.TensorFlowBackend.rq = _rq_
 tensornetwork.backends.tensorflow.tensorflow_backend.TensorFlowBackend.svd = _svd_tf
 
 
+def sparse_tensor_matmul(self: Tensor, other: Tensor) -> Tensor:
+    """
+    An implementation of matrix multiplication (@) for tf.SparseTensor.
+
+    This function is designed to be monkey-patched onto the tf.SparseTensor class.
+    It handles multiplication with a dense vector (rank-1 Tensor) by temporarily
+    promoting it to a matrix (rank-2 Tensor) for the underlying TensorFlow call.
+    """
+    # Ensure the 'other' tensor is of a compatible dtype
+    if not other.dtype.is_compatible_with(self.dtype):
+        other = tf.cast(other, self.dtype)
+
+    # tf.sparse.sparse_dense_matmul requires the dense tensor to be a 2D matrix.
+    # If we get a 1D vector, we need to reshape it.
+    is_vector = len(other.shape) == 1
+
+    if is_vector:
+        # Promote the vector to a column matrix [N] -> [N, 1]
+        other_matrix = tf.expand_dims(other, axis=1)
+    else:
+        other_matrix = other
+
+    # Perform the actual multiplication
+    result_matrix = tf.sparse.sparse_dense_matmul(self, other_matrix)
+
+    if is_vector:
+        # Demote the result matrix back to a vector [M, 1] -> [M]
+        return tf.squeeze(result_matrix, axis=1)
+    else:
+        return result_matrix
+
+
 class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):  # type: ignore
     """
     See the original backend API at `tensorflow backend
@@ -378,6 +416,8 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
             )
         tf = tensorflow
         tf.sparse.SparseTensor.__add__ = tf.sparse.add
+        tf.SparseTensor.__matmul__ = sparse_tensor_matmul
+
         self.minor = int(tf.__version__.split(".")[1])
         self.name = "tensorflow"
         logger = tf.get_logger()  # .setLevel('ERROR')
@@ -407,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)
 
@@ -490,12 +536,35 @@ 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)
 
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return tf.math.argmin(a, axis=axis)
 
+    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)
+
     def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         r = tf.unique_with_counts(a)
         order = tf.argsort(r.y)
@@ -504,6 +573,17 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
     def stack(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
         return tf.stack(a, axis=axis)
 
+    def clip(self, a: Tensor, a_min: Tensor, a_max: Tensor) -> Tensor:
+        return tf.clip_by_value(a, a_min, a_max)
+
+    def floor(self, a: Tensor) -> Tensor:
+        if a.dtype.is_integer:
+            return a
+        return tf.math.floor(a)
+
+    def floor_divide(self, x: Tensor, y: Tensor) -> Tensor:
+        return tf.math.floordiv(x, y)
+
     def concat(self, a: Sequence[Tensor], axis: int = 0) -> Tensor:
         return tf.concat(a, axis=axis)
 
@@ -678,7 +758,14 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
         sp_a: Tensor,
         b: Tensor,
     ) -> Tensor:
-        return tf.sparse.sparse_dense_matmul(sp_a, b)
+        is_vec = False
+        if len(b.shape) == 1:
+            b = self.reshape(b, [-1, 1])
+            is_vec = True
+        r = tf.sparse.sparse_dense_matmul(sp_a, b)
+        if is_vec:
+            return self.reshape(r, [-1])
+        return r
 
     def _densify(self) -> Tensor:
         @partial(self.jit, jit_compile=True)
@@ -712,7 +799,10 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
     def scan(
         self, f: Callable[[Tensor, Tensor], Tensor], xs: Tensor, init: Tensor
     ) -> Tensor:
-        return tf.scan(f, xs, init)[-1]
+        stacked_results = tf.scan(f, xs, init)
+        final_state = tf.nest.map_structure(lambda x: x[-1], stacked_results)
+        return final_state
+        # return tf.scan(f, xs, init)[-1]
 
     def device(self, a: Tensor) -> str:
         dev = a.device
@@ -864,7 +954,7 @@ class TensorFlowBackend(tensorflow_backend.TensorFlowBackend, ExtendedBackend):
         f: Callable[..., Any],
         static_argnums: Optional[Union[int, Sequence[int]]] = None,
         jit_compile: Optional[bool] = None,
-        **kws: Any
+        **kws: Any,
     ) -> Any:
         # static_argnums not supported in tf case, this is only for a consistent interface
         # for more on static_argnums in tf.function, see issue: https://github.com/tensorflow/tensorflow/issues/52193
@@ -1011,4 +1101,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)