how to assert with respect to traced values inside vmap?

[aldopareja]

Sorry if this has been asked before. This is what I want, this way I can make sure the function is never used under unexpected conditions.

def check_value(a):
    assert a.shape == (1,) #works
    assert a>0 #doesn't work

vmap(check_value)(jnp.ararge(10)-5)

[jakevdp]

You cannot write assertions with respect to traced values in JAX transforms (jit, vmap, etc.).

There's some more background on this in Common Gotchas: Python Control Flow – this section specifically discusses control flow, but assertions are basically control flow in disguise (assert a>0 is basically if a>0: raise AssertionError()).

The How To Think In JAX document is also some good background reading to understand why JAX behaves this way.

There may be workarounds depending on what you were trying to do originally: can you say more about what your goal was here? checkify and host_callback are both possible approaches to achieve similar results, depending on the context and how important performance is for your function.

[aldopareja]

Given the debugging tools that have spawned recently (e.g jax.debug.print), is there a way we can assert with respect to the values of traced arrays?

[jakevdp]

Sure, for example, you could use debug.callback to do a host-side assertion of the value:

import jax

def check_value(val):
  assert val > 0

def f(x):
  jax.debug.callback(check_value, x)
  return x

jax.vmap(f)(jax.numpy.array([1, 2, 3]))
# DeviceArray([1, 2, 3], dtype=int32)

jax.vmap(f)(jax.numpy.array([1, 2, -1]))
# AssertionError

This is probably not to be recommended in normal production code (the callback is done synchronously, so will likely make any real code much slower) but for debugging purposes it's a decent approach.

There's also jax.experimental.checkify which may eventually become a good approach for this kind of thing.

[aldopareja]

thank you so much!