Is there a way to bail out of JIT-ing?

[marius311]

The situation I'm in is I'm writing some library code which I want JIT-ed, but at some point I call a user provided function which may not be JIT-able. So ideally I want to write something like:

@jax.dontjit
def user_provided_unjittable_function(x):
    return np.exp(x)

@jax.jit
def my_library_function(x):
    # some stuff here I want jitted
    return user_provided_unjittable_function(x)

my_library_function(2)

where the user has indicated not to JIT their function with a hypothetical @jax.dontjit. Does anything like this exist in JAX?

I realize you can refactor my_library_function into something like:

def my_library_function(x):
    jax.jit(_jitted_parts)(...) # some stuff here I want jitted
    return user_provided_unjittable_function(x)

but in my real example it'd be much more complex and make the code less clear to do so.

[jakevdp]

You can use jax.disable_jit() as a context manager to run non-jitted versions of all code:

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  print(x)
  return x

x = jnp.arange(5)

f(x)
# prints Traced<ShapedArray(int32[5])>with<DynamicJaxprTrace(level=0/1)>

with jax.disable_jit():
  f(x)
  # prints [0 1 2 3 4]

As for calling non-jitted code within jitted code, it's not so easy: that would require on-device XLA code calling back to the Python runtime. You may be able to do something like this using host_callback, but I don't know of any example of exactly the scenario you have in mind.

[marius311]

Thanks, yea I think

As for calling non-jitted code within jitted code, it's not so easy

is exactly what the thing I'm after would entail, but hadn't realized its hard.

Based on your suggestion, I tried the following but couldn't get it working (jax 0.3.13), despite the docs sounding like it should (the docs say I should return a pytree, which I think I did, but the error message says it wants a PyTree(None)?)

def user_provided_unjittable_function(x):
    return np.exp(x) # intentionally use np not jnp here to simulate something unjittable

@jax.jit
def my_library_function(x):
    return jax.experimental.host_callback.call(user_provided_unjittable_function, x)

my_library_function(2)

ERROR:absl:Outside call <function _call.<locals>.<lambda> at 0x7f924457d310> threw exception Callback func <function _call.<locals>.<lambda> at 0x7f924457d310> should have returned a result with pytree PyTreeDef(None) but returned PyTreeDef(*).

Anyway, this isn't super critical but maybe I'm making a basic mistake using that?

[marius311]

Ahh read the docs a bit further to see you need to specify the result shape, so this works:

def user_provided_unjittable_function(x):
    return np.exp(x) # intentionally use np not jnp here to simulate something unjittable

@jax.jit
def my_library_function(x):
    return jax.experimental.host_callback.call(user_provided_unjittable_function, x, result_shape=x)

my_library_function(2.)

Pretty cool, will play with this option, thanks!