In order to work around long compile times and hidden re-compilations, a bunch of JIT compiled logic can be broken into a sequence of smaller functions.

One of the implications of this is that the control flow logic has to be host-side, since using jax.lax control flow functions triggers results in compilation of the target condition/body functions which is will effectively just compiling the original function. For example, if you were to break a function f into two functions f1 and f2 your code may look something like this with the loop running on the host:

state = compute_initial_state()

for i in range(n) :
    state = fn1(state)
    state = fn2(state)

print(state.result)

Making this work efficiently requires a few extra considerations:

1. Avoiding GPU-to-CPU transfer between calls to the JAX-compiled functions
2. Avoid re-allocations as much as possible through aggressive donation of arguments
3. Avoid hidden recompilation

This works but is fragile and it can be difficult to protect against regressions without burdensome performance testing. One thing that works great is that (3.) can be dealt with easily through AOT compilation, which will throw an exception whenever a regression is introduced that mistakenly changes the type of a return value (e.g. an array in state is initialized to [1] with a data type of int32 when it should have been initialized as [1.] with data type of float32).

For use cases like this, I think that it would be useful to have mechanisms that would protect against the following two cases:

1. A GPU-to-CPU transfer mistakenly took place within the critical loop, because host-side logic accessed a GPU variable.
2. An argument that should have been donated was not donated (e.g. could be introduced by changing of a method signature without updating the donation_argnames)

One approach would be to introducing new contexts that prevent these from happening inside a critical section of code. For example, with (2.) it could look something like the following (please ignore variable names, I did not put much thought into them):

state = compute_initial_state()
state = jax.device_put(state)

with jax.disable_transfer_to_cpu() :
    # Within this block, an exception will be thrown
    compiled_fn1 = compile(fn1, state)
    compiled_fn2 = compile(fn2, state)

    # This would throw an exception, since printing state.result would trigger
    # a transfer from GPU to CPU
    # print(state.result)

# This would not trigger an exception, since it is outside the protected block above
print(state.result)

Is this something that had been thought about or considered? I am not sure if this is an issue that other people have faced or not, and figured it was worth bringing up.