vmap and logical tests

[hovren]

I am trying to catch errors in a function which looks something like this:

def is_ok(x, xmin, xmax):
    return xmin <= x < xmax

def f(x):
    if not is_ok(x, -1., 1.):
        raise ValueError("Bad value")
    return x ** 2

Running this on single values work but using vmap does not

f(0.1)  # OK!
jax.vmap(f)(jnp.array([0.1, 0.2, 0.3]))  # Nope!

Where the latter raises the following error

---------------------------------------------------------------------------
FilteredStackTrace                        Traceback (most recent call last)
<ipython-input-55-49bbb7086bd9> in <module>
     11 
---> 12 jax.vmap(f)(jnp.array([0.1, 0.2, 0.3]))

<ipython-input-55-49bbb7086bd9> in f(x)
      4 def f(x):
----> 5     if not is_ok(x, -1., 1.):
      6         raise ValueError("Bad value")

<ipython-input-55-49bbb7086bd9> in is_ok(x, xmin, xmax)
      1 def is_ok(x, xmin, xmax):
----> 2     return xmin <= x < xmax
      3 

FilteredStackTrace: jax.core.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected.

The problem arose with the `bool` function. 



See https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error for more information.

Encountered tracer value: Traced<ShapedArray(bool[])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([ True,  True,  True], dtype=bool)
       batch_dim = 0

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

ConcretizationTypeError                   Traceback (most recent call last)
<ipython-input-55-49bbb7086bd9> in <module>
     10 f(jnp.array((0.4, )))
     11 
---> 12 jax.vmap(f)(jnp.array([0.1, 0.2, 0.3]))

~/miniconda3/envs/jaxplayground/lib/python3.8/site-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    137   def reraise_with_filtered_traceback(*args, **kwargs):
    138     try:
--> 139       return fun(*args, **kwargs)
    140     except Exception as e:
    141       if not is_under_reraiser(e):

~/miniconda3/envs/jaxplayground/lib/python3.8/site-packages/jax/api.py in batched_fun(*args, **kwargs)
   1220     in_axes_flat = flatten_axes("vmap in_axes", in_tree, (in_axes, 0), kws=True)
   1221     axis_size = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap", kws=True)
-> 1222     out_flat = batching.batch(
   1223         flat_fun, axis_name, axis_size, in_axes_flat,
   1224         lambda: flatten_axes("vmap out_axes", out_tree(), out_axes)

~/miniconda3/envs/jaxplayground/lib/python3.8/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    164 
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:
    168       # Some transformations yield from inside context managers, so we have to

<ipython-input-55-49bbb7086bd9> in f(x)
      3 
      4 def f(x):
----> 5     if not is_ok(x, -1., 1.):
      6         raise ValueError("Bad value")
      7     return x ** 2

<ipython-input-55-49bbb7086bd9> in is_ok(x, xmin, xmax)
      1 def is_ok(x, xmin, xmax):
----> 2     return xmin <= x < xmax
      3 
      4 def f(x):
      5     if not is_ok(x, -1., 1.):

~/miniconda3/envs/jaxplayground/lib/python3.8/site-packages/jax/core.py in __bool__(self)
    558   def __getitem__(self, idx): return self.aval._getitem(self, idx)
    559   def __nonzero__(self): return self.aval._nonzero(self)
--> 560   def __bool__(self): return self.aval._bool(self)
    561   def __int__(self): return self.aval._int(self)
    562   def __long__(self): return self.aval._long(self)

~/miniconda3/envs/jaxplayground/lib/python3.8/site-packages/jax/core.py in error(self, arg)
    953                       f"or `jnp.array(x, {fun.__name__})` instead.")
    954   def error(self, arg):
--> 955     raise_concretization_error(arg, fname_context)
    956   return error
    957 

~/miniconda3/envs/jaxplayground/lib/python3.8/site-packages/jax/core.py in raise_concretization_error(val, context)
    942          "See https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error for more information.

"
    943           f"Encountered tracer value: {val}")
--> 944   raise ConcretizationTypeError(msg)
    945 
    946 

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected.

The problem arose with the `bool` function. 



See https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error for more information.

Encountered tracer value: Traced<ShapedArray(bool[])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([ True,  True,  True], dtype=bool)
       batch_dim = 0

What is the correct way to handle this?

I should note that xmin and xmax above are dynamically computed in my actual code, but their values can be known at compile time (they only depend on other static values and a known shape).

[jakevdp]

This issue is discussed here: How to Think in JAX: Tracing and static variables.

In particular, when you apply a vmap transformation on code, the mapped variables are traced, and it is not possible to do control flow that depends on the values being mapped. So this statement:

    if not is_ok(x, -1., 1.):
        raise ValueError("Bad value")

is incompatible with vmap when x is a mapped value. You cannot raise an error conditional on mapped values.

The best course of action depends on what you're doing with the result. Can you say more about the application that this simple snippet is meant to represent?

[hovren]

I am trying to implement spline interpolation, and catching time points that extend outside of the spline interior at runtime would be nice.

Here is a short example where I am computing a dummy cost function. It works fine if I remove the call to is_valid_spline_time() in spline_evaluate():

def is_valid_spline_time(t, t0, dt, control_points):
    max_t = t0 + len(control_points) * dt
    return t0 <= t < max_t

def spline_evaluate(t, dt, t0, control_points, Mk):
    if not is_valid_spline_time(t, t0, dt, control_points):
        raise ValueError("Bad time")
        
    k = Mk.shape[0]
    t_local = t - t0
    i = jnp.floor_divide(t_local, dt).astype('int32').reshape(())
    u = (t_local - i * dt)/dt
    V = jax.lax.dynamic_slice_in_dim(control_points, i-k+1, k)
    U = u**jnp.arange(k)
    return U @ Mk @ V

M = jnp.array([[ 0.5,  0.5,  0. ],
       [-1. ,  1. ,  0. ],
       [ 0.5, -1. ,  0.5]])

rng = jax.random.PRNGKey(0)
key, rng = jax.random.split(rng)
control_points = jax.random.uniform(key, (10, 2), minval=0, maxval=10)
t0 = 0.0
dt = 1.0

max_t = t0 + len(control_points) * dt
test_ts = jnp.linspace(t0, max_t, endpoint=False, num=100)

key, rng = jax.random.split(rng)
points = jax.random.uniform(key, test_ts.shape, minval=0, maxval=10)


def loss(times, points, dt, t0, control_points, M):
    def loss_single(t, point):
        pt_spline = spline_evaluate(t, dt, t0, control_points, M)
        return jnp.linalg.norm(point - pt_spline)     
    
    losses = vmap(loss_single)(times, points)
    return jnp.sum(losses)
    

loss(test_ts, points, dt, t0, control_points, M)

[jakevdp]

JAX doesn't offer any mechanism to raise errors at runtime conditioned on traced values (see python control flow + JIT for details). Would returning NaN in these cases be suitable?

[hovren]

Using NaN could work, but I guess if the NaN is propagated through multiple functions then finding the piece of code that initially generated it can be harder.

Just to make sure I understand correctly: I have two problems here, right?

1. The kind of logical test I tried to write is not vmap- or jittable, and
2. Raising exceptions doesn't work anyway

So going with the NaN approach would solve problem (2) but I would still have to rewrite the test using jnp.where or jnp.lax.cond or something similar?

Regarding raising exceptions and error handling in general, is this something that is on the roadmap? Because being able to use the usual machinery of checking inputs and raising ValueError is very convenient for tracking down bugs.

[jakevdp]

The kind of logical test I tried to write is not vmap- or jittable

Sort of: the kind of logical test you tried to write is jittable, it's just that the result of the test must be changing other traced values, not affecting Python control flow (which is only possible at compile time). For example, you cannot do this:

@jit
def f(x):
  if x.sum() < 0:
    return jnp.zeros_like(x)
  else:
    return x

but you can do this:

@jit
def f(x):
  return jnp.where(x.sum() < 0, jnp.zeros_like(x), x)

Raising exceptions doesn't work anyway

Raising exceptions works at compile time (conditioned on static values), but not at runtime (conditioned on content of arrays), because XLA running on accelerators like GPUs don't really provide any way to propagate runtime errors back to the host.

is this something that is on the roadmap?

I don't know of any efforts to introduce runtime errors into XLA computations, but you might be able to do something like this via host_callback, though I imagine having a callback for each value in your array would likely have performance implications.

Another option would be to have the check in a non-jitted wrapper; something like this:

def myfunc(x):
  if jnp.any(x < 0):
    raise ValueError(...)
  return _myfunc(x)

@jit
def _myfunc(x):
  ...