I think this syntax
start = time.time()
result2 = jnp.cos(X)
jax.device_get(result2) # Force execution completion
second_time = time.time() - start
is nicer than .block_until_ready() for timing.
Could someone please test this and confirm that it's a suitable alternative. If yes we'll change throughout.
CC @mmcky -- perhaps you could coordinate with @HumphreyYang or Bishmay?
I also prefer that we move away from %time and %%time because they cause problems with jupytext conversion.
I think this syntax
is nicer than
.block_until_ready()for timing.Could someone please test this and confirm that it's a suitable alternative. If yes we'll change throughout.
CC @mmcky -- perhaps you could coordinate with @HumphreyYang or Bishmay?
I also prefer that we move away from %time and %%time because they cause problems with jupytext conversion.