Avoid jitting jax functions inside a jitted function

[AlbertMitjans]

Is there any way to disable the jitting of jax functions inside a jitted function?

If I call the dummy function defined below, an error will be raised, because jnp.iscomplex(x) returns a tracer.

But x is fixed, thus I'd want jnp.iscomplex(x) to return False.

import jax
import jax.numpy as jnp

x = jnp.array(3)

@jax.jit
def dummy():
  if jnp.iscomplex(x):
    print("Is complex!")

[jakevdp]

In general you cannot avoid jit-compiling jax functions within a JIT function. The problem is that the return value of jnp.iscomplex depends on the values in the input, and those values are not known at trace-time (when Python control flow executes) but only at runtime. There are three possible approaches to what you want, depending on the application:

Use static properties in control flow

You can do python control flow that branches on static rather than dynamic attributes. Examples of static attributes are the shape and dtype of an array. For example:

@jax.jit
def dummy():
  if jnp.issubdtype(x.dtype, jnp.complexfloating):
    print("Is complex!")

Use numpy for static computations

If you'd like to do static computations on static values, you can do so with numpy rather than jax.numpy functions:

import numpy as np

x = np.array(3)  # note: numpy not jax.numpy

@jax.jit
def dummy():
  if np.iscomplex(x):  # note: numpy not jax.numpy
    print("Is complex!")

External Callbacks

If you're doing something more sophisticated that you would like executed at runtime, you can do this with external callbacks, although this typically comes with performance penalties due to the data transfer and synchronization between host and device. You can read more about that here: https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html

[soraros]

Could also suggest np.iscomplexobj.

[jakevdp]

Good point! It's basically a simpler way writing option 1: https://github.com/google/jax/blob/5cfd15bd19cefe01e9437f6de6b45d23cb7fc1a9/jax/_src/numpy/lax_numpy.py#L142-L147

[soraros]

@jakevdp Out of curiosity, why can't we simply define iscomplexobj as np.iscomplexobj in lax_numpy.py?

[jakevdp]

It doesn't work correctly with jit in all situations, because it calls np.asarray on inputs. e.g.

>>> jax.jit(np.iscomplexobj)([1])
TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
The error occurred while tracing the function iscomplexobj at /Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.8/site-packages/numpy/lib/type_check.py:303 for jit. This concrete value was not available in Python because it depends on the value of the argument 'x'.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

[soraros]

Ah, that makes sense. I somehow forget #11581. Thanks!