Debugging "No constant handler" errors

[DanPuzzuoli]

Hi all,

With the 0.4.4 update an old-ish test case in a package I work on is now failing and raising the following error:

TypeError: No constant handler for type: <class 'jax.interpreters.partial_eval.DynamicJaxprTracer'>

I haven't encountered this kind of error before so I'm not totally sure where to start. The full error can be viewed in this failed workflow - there isn't any information about where exactly in my code base this is being raised. (Note that the test that's failing is unrelated to the code in the linked PR.)

I'm going to try to dig deeper on my own, but wanted to ask here in case anything comes to mind for what I should look out for. Unfortunately this is testing some of the most complicated functionality in this package (utilizing multiple layers of transformations), and some strange edge cases have popped up here during JAX updates before.

Thanks!

[DanPuzzuoli]

One thing to note: the line in my test that raises the error is of the form:

jit(grad(f))(1.0)

(f takes a scalar input). If I change this line in the test to just:

grad(f)(1.0)

it passes. So there seems to be some problem with jitting the grad of f.

[jakevdp]

Can you provide a minimal reproduction of the issue? I don't see the error when I run this in a simple case:

In [1]: from jax import jit, grad
In [2]: f = lambda x: x
In [3]: jit(grad(f))(1.0)
Out[3]: Array(1., dtype=float32, weak_type=True)

[DanPuzzuoli]

This doesn't reproduce exactly the same error, but this is a small case I'd used when this test failed after a previous JAX update a while ago (link to that discussion). It seems to have all the same characteristics as are described in that discussion: some particular combination of vmap, switch, odeint, jit and grad causes an error, but if the example is modified by removing any one of these things it works.

Update: nevermind, it is the same error, but it was just printing more afterward so I thought it was different.

# setting this makes things work
#import os
#os.environ['JAX_JIT_PJIT_API_MERGE'] = '0'

import numpy as np
import jax.numpy as jnp
from jax import vmap, jit, grad
from jax.lax import switch
from jax.experimental.ode import odeint

def A0(t):
    return 2.

def A1(a, t):
    return a**2

y0 = np.random.rand(2)
T = np.pi * 1.232

def test_func(a):
    
    eval_list = [A0, lambda t: A1(a, t)]
    
    def single_eval(idx, t):
        return switch(idx, eval_list, t)
    multiple_eval = vmap(single_eval, in_axes=(0, None))
    idx_list = jnp.array([0, 1])
    rhs = lambda y, t: multiple_eval(idx_list, t) * y
    
    # using this version of the RHS works so the vmap seems to be necessary
    #rhs = lambda y, t: jnp.array([single_eval(0, t), single_eval(1, t)]) * y
    
    # using this version also seems to work
    # so issue may be with vmap + switch in this context?
    #rhs = lambda y, t: vmap(lambda x: a * x)(y)
    
    # returning this instead also works, so odeint is necessary?
    #return multiple_eval(idx_list, 1.) * jnp.ones((2, 2), dtype=complex)
    
    out = odeint(
        rhs,
        y0=y0,
        t=jnp.array([0, T], dtype=float),
        atol=1e-13,
        rtol=1e-13
    )
    return out

jit(grad(lambda a: test_func(a)[-1][1].real))(1.)

# evaluation without jitting works
#grad(lambda a: test_func(a)[-1][1].real)(1.)

[DanPuzzuoli]

Through binary searching on main commit history, using jaxlib==0.4.4, the error appears in the commit merging jit and pjit. Not sure what's going on, but if it's changing a flag to use code that was developed in other PRs this may not be the most informative.

[yashk2810]

Do you have the latest jaxlib installed? You need jaxlib 0.4.4 with jax 0.4.4

[DanPuzzuoli]

Yes - the behaviour of the minimal example above is with both jaxlib and jax being 0.4.4.

[mattjj]

I suspect this issue is with the odeint custom_vjp's use of nondiff_argnums.

[yashk2810]

Uh oh! custom_vjp strikes again! Nice catch Matt! The func is the problem which might contain tracers?

[mattjj]

Yes, though the other args could be problematic too. I can take this. We should probably make it a bug issue and assign me (and maybe Roy too because it's related to another bug we're on).

[DanPuzzuoli]

@mattjj just wanted nudge to see if there are any updates with this, as I see that in 0.4.7 it is no longer possible to revert to the old JIT pathway with:

import os
os.environ['JAX_JIT_PJIT_API_MERGE'] = '0'