Detect when a new JAX frame starts

[dionhaefner]

For mpi4jax, our API would get a lot cleaner if we could keep track of local state within each jit block to inject tokens into our custom calls. But for this I need to detect when a new frame starts to clear the local state.

I can make this work with something like this:

current_token = None

def state_wrapper(func):
    def wrapped(*args, **kwargs):
        global current_token
        if current_token is None:
            current_token = jax.lax.create_token()
        try:
            out = func(*args, **kwargs, token=current_token)
        # this is the part that feels icky
        except jax.core.UnexpectedTracerError:
            current_token = jax.lax.create_token()
            out = func(*args, **kwargs, token=current_token)
        current_token = out[-1]
        return out[:-1]
    return wrapped

Then we can use it like this:

# this is what we do so far
@jax.jit
def old(arr):
    token = jax.lax.create_token()
    arr, token = custom_call(arr, token=token)
    arr, token = custom_call(arr, token=token)
    return arr

# no need to keep track of tokens manually anymore
custom_call = state_wrapper(custom_call)

# this is what it could look like
@jax.jit
def new(arr):
    arr = custom_call(arr)
    arr = custom_call(arr)
    return arr   

Is there any way to detect that we are now in a different frame as the one that current_token originated from that doesn't rely on catching UnexpectedTracerError?