Are there jit compile-time guards?

[tillahoffmann]

Are there jit compile-time guards to conditionally execute code akin to #ifdef in C++? For context, I'd like to implement value assertions, which are not allowed in jit-ed code, for debugging. Something like this would address my use case.

>>> def custom_log(x):
...     if not jax.is_jitting:  # Or a better name ...
...         np.testing.assert_array_less(0, x)
...     return jnp.log(x)

>>> jax.jit(custom_log)(jnp.zeros(1))
Array([-inf], dtype=float32)

>>> custom_log(jnp.zeros(1))
AssertionError: 
Arrays are not less-ordered

Mismatched elements: 1 [/](http://127.0.0.1:8888/) 1 (100%)
Max absolute difference: 0.
Max relative difference: inf
 x: array(0)
 y: array([0.], dtype=float32)

Currently, I just comment out all my assertions before jitting and then uncomment them for more debugging. Is there a better approach?

[jakevdp]

No, we deliberately avoid providing any API along those lines, because as soon as you write code that operates differently at trace-time and run-time, you break JAX's tracing model and open the door to difficult to track bugs (e.g. grad would return the gradient of the trace-time code, not the runtime code, which could break an ML model in hard-to-debug ways).

Even if it were possible, the "compile-time only" assertion approach wouldn't even work correctly, because we don't recompile on subsequent function calls. For example:

jax.jit(custom_log)(jnp.array([1.0]))  # passes the assertion
jax.jit(custom_log)(jnp.array([0.0]))  # hits JIT cache, so compile-time assertion doesn't run

It sounds like what you're after is not compile-time assertions, but rather runtime value assertions.

There are essentially two supported ways to do that right now: jax.experimental.checkify, or jax.debug.callback.

For what it's worth, JAX itself avoids this kind of runtime assertion entirely, instead returning invalid values like NaN, because checkify is somewhat inconvenient to use, and runtime callbacks can have performance implications, particularly on accelerators. There's some discussion of this in https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#out-of-bounds-indexing

[tillahoffmann]

Thank you for the fast reply. I was thinking of conditional runtime assertions that are only executed when code is not jitted. In the example above it would be

jax.jit(custom_log)(jnp.array([1.0]))  # doesn't run the assertion
jax.jit(custom_log)(jnp.array([0.0]))  # doesn't run the assertion
custom_log(jnp.array([1.0]))  # checks and passes
custom_log(jnp.array([0.0]))  # checks and fails

I appreciate the danger of possibly difficult bugs to hunt down, I hadn't thought of that. This is a sort of hacky implementation of what I had in mind.

import functools
import inspect
import jax
from jax import numpy as jnp
import numpy as np


@functools.wraps(jax.jit)
def condjit(
    fun,
    in_shardings=None,
    out_shardings=None,
    static_argnums=None,
    static_argnames=None,
    donate_argnums=None,
    donate_argnames=None,
    keep_unused=False,
    device=None,
    backend=None,
    inline=False,
    abstracted_axes=None,
):
    # Inject static `jitted` parameter if the signature accepts it.
    signature = inspect.signature(fun)
    if "jitted" in signature.parameters:
        static_argnames = static_argnames or []
        static_argnames.append("jitted")
        fun = functools.partial(fun, jitted=True)
    return jax.jit(
        fun, in_shardings, out_shardings, static_argnums, static_argnames, 
        donate_argnums, donate_argnames, keep_unused, device, backend, inline, 
        abstracted_axes
    )


def custom_log(x, jitted=False):
    if not jitted:
        np.testing.assert_array_less(0, x)
    return jnp.log(x)


condjit(custom_log)(jnp.array([1.0]))  # doesn't run the assertion
condjit(custom_log)(jnp.array([0.0]))  # doesn't run the assertion
custom_log(jnp.array([1.0]))  # checks and passes
custom_log(jnp.array([0.0]))  # checks and fails

[jakevdp]

It's certainly possible to write code that does something like what you have in mind using a number of mechanisms, but we do not recommend it.

[tillahoffmann]

Okay, thank you for the fast replies. Have a good weekend!

[jakevdp]

Now that I think about it, there is one way to safely implement this: a custom primitive that does the assertion in the impl rule, and does nothing in the lowering rule (and the batching rule, autodiff rule, etc.). Something like this:

import jax
import jax.numpy as jnp
from jax import core
from jax.interpreters import ad, batching, mlir

def assert_if_not_traced(x, message):
  return assert_if_not_traced_p.bind(x, message=message)

assert_if_not_traced_p = core.Primitive('assert_if_not_traced')
assert_if_not_traced_p.def_abstract_eval(lambda x, **_: x)
@assert_if_not_traced_p.def_impl
def _assert_if_not_traced_impl(x, *, message):
  assert x, message
  return x
batching.defvectorized(assert_if_not_traced_p)    
ad.deflinear2(assert_if_not_traced_p, lambda t, _: [t])
mlir.register_lowering(assert_if_not_traced_p, mlir.lower_fun(lambda x, **_: x))

Then you can use it this way:

def f(x):
  assert_if_not_traced(jnp.all(x > 0), "x must be positive")
  return x  # return value needed to demonstrate `grad` below.

f(1)  # no error (valid value)
jax.jit(f)(0)  # no error (traced)
jax.jit(jax.grad(f))(0.0)  # no error (traced)
jax.jit(jax.vmap(f))(jnp.zeros(1))  # no error (traced)

jax.grad(f)(0.0)  # AssertionError: x must be positive
jax.vmap(f)(jnp.zeros(1))  # AssertionError: x must be positive
f(0)  # AssertionError: x must be positive

[tillahoffmann]

Interesting, thank you for the example. I'll have to read more about the internals to get my head around the code.

[patrick-kidger]

Two other reference points you might find useful here:

• eqx.internal.{nonbatchable, nondifferentiable, ...}. These are custom primitives working in a very similar way to Jake's example at the end -- in this case they raise errors when vmap'd etc.

• eqx.error_if, which offers runtime errors (using callbacks, also mentioned by Jake earlier) under JIT. These already support a couple of options (e.g. opening up a breakpoint if they fail); we could probably add a flag to these that makes them only trigger when ran without JIT.

FWIW I think Jake's final comment is exactly what you want for the use-case you described, but I wanted to reference these as useful tools to have in your toolbox when tackling problems of this type.