Repository: scipy/scipy
File: benchmarks/benchmarks/common.py (Python)

Code before:

        elif backend == "dask.array":
            import dask

            def wrapper(*args, **kwargs):
                res = func(*args, **kwargs)
                return dask.compute(res)[0]

            self.func = wrapper

        elif backend == "cupy":
            import cupy
            # The default stream is non-blocking.
            # As of CuPy 13.4.1, explicit non-blocking streams
            # are substantially slower.
            # cupy.cuda.Stream(non_blocking=True).use()

            def wrapper(*args, **kwargs):
                res = func(*args, **kwargs)
                cupy.cuda.get_current_stream().synchronize()
                return res

            self.func = wrapper

        else:
            assert backend in ("numpy", "array_api_strict", "cupy")

    def synchronize(self, *arrays):
        """Wait until the given arrays have finished generating and return a
        synchronized instance of them.
        You need to call this on all arrays that your setup() function creates.
        """
        if self.backend == "dask.array":
            import dask

            arrays = dask.persist(*arrays)
        elif self.backend in ("jax.numpy:cpu", "jax.numpy:cuda"):
            import jax

            jax.block_until_ready(arrays)
        elif self.backend == "torch:cuda":
            import torch

            torch.cuda.synchronize()
        elif self.backend == "cupy":
            import cupy

            cupy.cuda.get_current_stream().synchronize()
        else:
            assert self.backend in ("numpy", "array_api_strict", "torch:cpu")


Review comment (suggestion):
```suggestion
            assert backend in ("numpy", "array_api_strict")
```

Code after:
                res = func(*args, **kwargs)
                jax.block_until_ready(res)
                return res

            self.func = wrapper

        elif backend == "dask.array":
            import dask

            def wrapper(*args, **kwargs):
                res = func(*args, **kwargs)
                return dask.compute(res)[0]

            self.func = wrapper

        elif backend == "cupy":
            import cupy
            # The default stream is non-blocking.
            # As of CuPy 13.4.1, explicit non-blocking streams
            # are substantially slower.
            # cupy.cuda.Stream(non_blocking=True).use()

            def wrapper(*args, **kwargs):
                res = func(*args, **kwargs)
                cupy.cuda.get_current_stream().synchronize()
                return res

            self.func = wrapper

        else:
            assert backend in ("numpy", "array_api_strict")

    def synchronize(self, *arrays):
        """Wait until the given arrays have finished generating and return a
        synchronized instance of them.
        You need to call this on all arrays that your setup() function creates.
        """
        if self.backend == "dask.array":
            import dask

            arrays = dask.persist(*arrays)
        elif self.backend in ("jax.numpy:cpu", "jax.numpy:cuda"):
            import jax

            jax.block_until_ready(arrays)
        elif self.backend == "torch:cuda":
            import torch

            torch.cuda.synchronize()
        elif self.backend == "cupy":
            import cupy