Omnistaging and interactive programming?

[tshadley]

Python lends itself well to an interactive style of programming:

1. set debugger breakpoints (pdb.set_trace() for example) at a point in your new code where you're not exactly sure what code is needed.
2. Execute the code, enter the breakpoint, and work out exactly what code is needed using python/NumPy/JAX commands in real time on the actual inputs.
3. Add new code and repeat the process with a new breakpoint.

This approach works well when the code can be executed quickly and overall iterations short. I use it all the time for new designs, getting things right with the first set of inputs.

However, with JAX omnistaging, it seems that most JAX/FLAX code is compiled to XLA (regardless of whether JIT is disabled or not), creating lengthy delays and making this interactive style painfully slow.

Is there a way to speed up, or avoid entirely, compiling to XLA?

I have heard, in theory, replacing jax.numpy calls with numpy is a workable solution-- most of an algorithm/FLAX module can be rapidly worked out using numpy and then converted to JAX. However, in practice, I know of no way of doing this except tediously, call by call.

Appreciate any thoughts and ideas on this.

Best to all

[jakevdp]

You can use the disable_jit context manager:

with jax.disable_jit():
  do_something_without_jit()

Or to disable JIT globally, you can use the jax_disable_jit flag, e.g. with

from jax import config
config.update('jax_disable_jit', True)

More at https://jax.readthedocs.io/en/latest/jax.html#jax.disable_jit

[tshadley]

@jakevdp, thanks for the reply.

I am disabling JIT (using config.update('jax_disable_jit', True) but my understanding and observation is that this has no effect on omnistaging and subsequent slow compiling time.

I'm using the environment variable JAX_LOG_COMPILES=1 to verify that delays are due to XLA compiling.

Here's an example. export JAX_LOG_COMPILES=1 before running.

import jax
import os
jax.config.update('jax_disable_jit', True)
jax.config.update('jax_platform_name', 'cpu')
rng = jax.random.PRNGKey(0)

I get (jax 0.2.25):
WARNING:absl:Compiling prim_fun (140434808386688) for args (ShapedArray(int32[]),).
WARNING:absl:Compiling prim_fun (140434810830400) for args (ShapedArray(int32[]), ShapedArray(int32[])).
WARNING:absl:Compiling prim_fun (140434808389056) for args (ShapedArray(int32[]),).
WARNING:absl:Compiling prim_fun (140434973779008) for args (ShapedArray(uint32[]),).
WARNING:absl:Compiling prim_fun (140434805740096) for args (ShapedArray(uint32[]),).
WARNING:absl:Compiling prim_fun (140434805604608) for args (ShapedArray(int32[]), ShapedArray(int32[])).
WARNING:absl:Compiling prim_fun (140434805725952) for args (ShapedArray(uint32[1]), ShapedArray(uint32[1])).