run lax.while_loop for a certain amount of time

[jeremiecoullon]

Is it possible to run a lax.while_loop until a certain amount of time has elapsed?

An example of what I mean: define cond_fun in the while_loop documentation as: lambda start_time: time.time() - start_time < 10

I've noticed that calling time.time() in a jitted function returns the current time only when the function is compiled but not anytime after that.

I would imagine that this isn't possible but I just thought I'd check!

[seanmor5]

I know it's been awhile, but I thought this was interesting. Here's an approach using the experimental Host Callback API:

from jax import lax, jit, partial
from jax.experimental import host_callback as hcb
import numpy as np
import time

# Call expects a value to be passed, so we ignore first arg
def stop_after_time(_, time_limit=None, start_time=None):
  cur_time = time.time()
  return np.array((cur_time - start_time) < time_limit)

@partial(jit, static_argnums=(1, 2,))
def _cond_fn(input, start_time=None, time_limit=None):
  # Partial application to avoid HCB Cache
  hcb_fun = partial(stop_after_time, time_limit=time_limit, start_time=start_time)
  # Send a blank value
  return hcb.call(hcb_fun, (), result_shape=jax.ShapeDtypeStruct((), np.bool))

@jit
def body_fn(acc):
  return acc + 1

@partial(jit, static_argnums=(1, 2,))
def sum_until_time(init_value, start_time, time_limit):
  cond_fn = partial(_cond_fn, start_time=start_time, time_limit=time_limit)
  return lax.while_loop(cond_fn, body_fn, init_value)

sum_until_time(0, time.time(), 10)

[seanmor5]

A note for anybody reading would be that depending on the time it takes to compile the computations, the loop might not run for any iterations so this probably isn't a good idea in general. You could probably include a "warm-up" time or something else to get around this, but I don't think there's any way to guarantee the loop runs for n seconds not including compilation time

[jeremiecoullon]

This is really great thanks!! I didn't realise there was a call() function in host_callback; nice!

About the compilation time: one solution is to run _cond_fn() with the given start and end time first to compile it. Then you can run lax.while_loop.

[jeremiecoullon]

Actually my last point about the compilation time is wrong: this wouldn't work as the start_time would still be before compilation has started.
Perhaps a way around this would be to have the start_time not be a static argument so you can pass in a new start time once _cond_fn is compiled?