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

Code before:
                torch.empty(0, device=device)
            except RuntimeError:
                raise SkipNotImplemented(f"{device=} not available")
            torch.set_default_device(device)

            if device == "cuda":
                def wrapper(*args, **kwargs):
                    res = func(*args, **kwargs)
                    torch.cuda.synchronize()
                    return res

                self.func = wrapper

        elif backend == "jax.numpy":
            import jax

            jax.config.update("jax_enable_x64", True)
            try:
                jax_device = jax.devices(device)[0]
            except RuntimeError:
                raise SkipNotImplemented(f"{device=} not available")
            jax.config.update("jax_default_device", jax_device)

            if capabilities["jax_jit"]:
                func = jax.jit(func, static_argnames=static_argnames,
                                static_argnums=static_argnums)
                self.warmup = True

            def wrapper(*args, **kwargs):
                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()


Review comment (suggestion):
```suggestion
                               static_argnums=static_argnums)
```

Code after:
        if backend == "torch":
            import torch

            torch.set_default_dtype(torch.float64)
            try:
                torch.empty(0, device=device)
            except (RuntimeError, AssertionError):
                raise SkipNotImplemented(f"{device=} not available")
            torch.set_default_device(device)

            if device == "cuda":
                def wrapper(*args, **kwargs):
                    res = func(*args, **kwargs)
                    torch.cuda.synchronize()
                    return res

                self.func = wrapper

        elif backend == "jax.numpy":
            import jax

            jax.config.update("jax_enable_x64", True)
            try:
                jax_device = jax.devices(device)[0]
            except RuntimeError:
                raise SkipNotImplemented(f"{device=} not available")
            jax.config.update("jax_default_device", jax_device)

            if capabilities["jax_jit"]:
                func = jax.jit(func, static_argnames=static_argnames,
                               static_argnums=static_argnums)
                self.warmup = True

            def wrapper(*args, **kwargs):
                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