Explicitly wait until computations finish in the verbose test. Should be NFC since currently `jax.jit` is sync on CPU.

google / jaxopt

Hardware accelerated, batchable and differentiable optimizers in JAX.

https://jaxopt.github.io

Apache License 2.0

903 stars 62 forks source link

Closed copybara-service[bot] closed 3 months ago

copybara-service[bot] commented 3 months ago

Explicitly wait until computations finish in the verbose test. Should be NFC since currently jax.jit is sync on CPU.

Context: we plan to make jax.jit async on CPU backend for expensive computations, which will fail this test without this change.