add searchsorted backend method

Author: refraction-ray

CHANGELOG.md

@@ -18,6 +18,8 @@
 
 - Add circuit `from_qsim_file` method to load Google random circuit structure
 
+- Add `searchsorted` method for backend
+
 ### Changed
 
 - The inner mechanism for `sample_expectation_ps` is changed to sample representation from count representation for a fast speed

tensorcircuit/backends/abstract_backend.py

@@ -746,6 +746,26 @@ def solve(self: Any, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
             "Backend '{}' has not implemented `solve`.".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.
+
+        :param a: input array sorted in ascending order
+        :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.
+            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
+        :return: Array of insertion points with the same shape as v, or an integer if v is a scalar.
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `searchsorted`.".format(self.name)
+        )
+
     def tree_map(self: Any, f: Callable[..., Any], *pytrees: Any) -> Any:
         """
         Return the new tree map with multiple arg function ``f`` through pytrees.

tensorcircuit/backends/jax_backend.py

@@ -389,6 +389,13 @@ def is_tensor(self, a: Any) -> bool:
     def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:  # type: ignore
         return jsp.linalg.solve(A, b, assume_a)
 
+    def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
+        if not self.is_tensor(a):
+            a = self.convert_to_tensor(a)
+        if not self.is_tensor(v):
+            v = self.convert_to_tensor(v)
+        return jnp.searchsorted(a, v, side)
+
     def tree_map(self, f: Callable[..., Any], *pytrees: Any) -> Any:
         return libjax.tree_map(f, *pytrees)
 

tensorcircuit/backends/numpy_backend.py

@@ -223,6 +223,9 @@ def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:  # type:
         # https://stackoverflow.com/questions/44672029/difference-between-numpy-linalg-solve-and-numpy-linalg-lu-solve/44710451
         return solve(A, b, assume_a)
 
+    def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
+        return np.searchsorted(a, v, side=side)  # type: ignore
+
     def set_random_state(
         self, seed: Optional[int] = None, get_only: bool = False
     ) -> Any:

tensorcircuit/backends/pytorch_backend.py

@@ -389,6 +389,13 @@ def left_shift(self, x: Tensor, y: Tensor) -> Tensor:
     def solve(self, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
         return torchlib.linalg.solve(A, b)
 
+    def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
+        if not self.is_tensor(a):
+            a = self.convert_to_tensor(a)
+        if not self.is_tensor(v):
+            v = self.convert_to_tensor(v)
+        return torchlib.searchsorted(a, v, side=side)
+
     def reverse(self, a: Tensor) -> Tensor:
         return torchlib.flip(a, dims=(-1,))
 

tensorcircuit/backends/tensorflow_backend.py

@@ -434,6 +434,9 @@ def solve(self, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
             return self.reshape(x, x.shape[:-1])
         return x
 
+    def searchsorted(self, a: Tensor, v: Tensor, side: str = "left") -> Tensor:
+        return tf.searchsorted(a, v, side)
+
     def from_dlpack(self, a: Any) -> Tensor:
         return tf.experimental.dlpack.from_dlpack(a)
 

tensorcircuit/experimental.py

@@ -226,7 +226,6 @@ def parameter_shift_grad(
     :return: the grad function
     :rtype: Callable[..., Tensor]
     """
-    # TODO(@refraction-ray): finite shot sample_expectation_ps not supported well for now
     if jit is True:
         f = backend.jit(f)
 
@@ -283,7 +282,6 @@ def parameter_shift_grad_v2(
     :return: the grad function
     :rtype: Callable[..., Tensor]
     """
-    # TODO(@refraction-ray): finite shot sample_expectation_ps not supported well for now
     if jit is True:
         f = backend.jit(f)
 

tests/test_backends.py

@@ -279,6 +279,10 @@ def test_backend_methods_2(backend):
         np.array([1, 3]),
     )
     assert tc.backend.dtype(tc.backend.ones([])) == "complex64"
+    edges = [-1, 3.3, 9.1, 10.0]
+    values = tc.backend.convert_to_tensor(np.array([0.0, 4.1, 12.0], dtype=np.float32))
+    r = tc.backend.numpy(tc.backend.searchsorted(edges, values))
+    np.testing.assert_allclose(r, np.array([1, 2, 4]))
 
 
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])