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

Code before:
                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")

        return arrays


def is_xslow():
    try:
        return int(os.environ.get('SCIPY_XSLOW', '0'))
    except ValueError:
        return False


class LimitedParamBenchmark(Benchmark):
    """
    Limits parameter combinations to `max_number` choices, chosen
    pseudo-randomly with fixed seed.
    Raises NotImplementedError (skip) if not in active set.
    """
    num_param_combinations = 0

    def setup(self, *args, **kwargs):

Review comment (suggestion):
Today this is unnecessary on `jax.numpy:cpu`, but it's not spelled out explicitly in the documentation and it's perfectly reasonable for JAX to kick off a worker thread in the future - or maybe even lean on OpenCL.

Code after:
            # 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

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

        return arrays[0] if len(arrays) == 1 else arrays


def is_xslow():
    try:
        return int(os.environ.get('SCIPY_XSLOW', '0'))
    except ValueError:
        return False


class LimitedParamBenchmark(Benchmark):
    """
    Limits parameter combinations to `max_number` choices, chosen
    pseudo-randomly with fixed seed.