How to detect if within jit?

[antalszava]

Hi everyone, is there a way to detect from within a function, if it's being jitted e.g., to create a function that works differently when jitted?

In the following example, I'm curious if we can have jitting be dependent on whether or not jax.jit is being called on func:

import jax
from jax import numpy as jnp
from jax.experimental.host_callback import call

def func(param):
    
    jitting = True # detect here somehow if we're jitting
    if jitting:
        shape = jax.ShapeDtypeStruct((), jnp.float32)
        return call(jnp.sin, param, result_shape=shape)
    return jnp.sin(param)

func(jnp.pi) # This would not use host_callback.call, jitting=False
jax.jit(func)(jnp.pi) # This would use host_callback.call, jitting=True

[jakevdp]

I don't think there's any public API for this, but you can detect it based on the side-effects, e.g. the fact that array creation results in a tracer within JIT and other transforms:

import jax.numpy as jnp
from jax import core, jit

def f():
  is_jit = isinstance(jnp.array(0), core.Tracer)
  return is_jit

print(f())
# False

print(jit(f)())
# True

(Note however this will return True not just for JIT, but any JAX transform, including grad, vmap, etc.)

That said, forking a function's behavior based on whether it is being called within a transform is probably not a good idea: it's likely to have unintended side-effects, particularly when it comes to things like jit-invariance and composability of transforms.

[antalszava]

Thank you for the fast response! 🙂 Indeed, it seems like when jitting, arrays will be instances of core.Tracer and in specific jax.interpreters.partial_eval.DynamicJaxprTracer, although that seems to be the case for example for pmap too.

That said, forking a function's behavior based on whether it is being called within a transform is probably not a good idea: it's likely to have unintended side-effects, particularly when it comes to things like jit-invariance and composability of transforms.

Does jit-invariance refer to the concept that the behaviour (e.g., outputs, errors raised) of a function doesn't change just by jitting it? What could be some implications to the composability of transforms?

[jakevdp]

Does jit-invariance refer to the concept that the behaviour (e.g., outputs, errors raised) of a function doesn't change just by jitting it?

Yes. This is an explicit goal of every function provided by JAX, and we work hard to maintain it. If you let implementations diverge based on whether they are traced, that hard work becomes your own responsibility!

What could be some implications to the composability of transforms?

Given your snippet above, I'd be worried for example that vmap(jit(f)) could end up having different behavior than jit(vmap(f)), given that under jit you make different assumptions about the shape of the input & output.

[wookayin]

Is there any alternative way to detect whether a function runs under jit but not under vmap? The reason I would need this is to assert unexpected side-effects would not happen:

        assert not isinstance(x, jax.core.Tracer), \
          "This function should not be JIT compiled due to a stateful side effect."

where vmap transformation of the function would be fine.

[jakevdp]

I don't know of any public API that allows detecting that. If you're depending on side-effects, I'd suggest disallowing traced code altogether, as vmap or grad will probably lead side effects to happen in unexpected ways as well. Alternatively, you could use something like io_callback (once it lands) to execute effectful code.

[cdagher]

I'm interested in this as well. I would like a way to warn a user that a function is not transformable, and was thinking of creating a decorator for this. Is it possible?