Abstract tracer value encountered

[helange23]

Hello,

I am fairly new to jax and I am encountering the following error:
ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected (in bool).

I have read https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error and I think I understand what causes the error but not how to remedy it.

def moeller_trumbore(p1, p2, q1, q2, q3):

    EPS = 1E-6
    
    vertex0 = q1
    vertex1 = q2
    vertex2 = q3
    
    edge1 = vertex1 - vertex0
    edge2 = vertex2 - vertex0
    
    h = jnp.cross(p2,edge2)
    a = jnp.dot(edge1,h)
    
    if jnp.abs(a) < EPS:
        return (False, 'a', 0)
    
    f = 1./a
    s = p1 - vertex0
    
    u = f*jnp.dot(s,h)
    
    if (u<0) or (u>1.0):
        return (False, 'u', 0)
    
    q = jnp.cross(s,edge1)
    v = f*jnp.dot(p2,q)
    
    if (v<0) or (u + v > 1.0):
        return (False, 'v', 0)
    
    t = f*jnp.dot(edge2,q)
    
    if (t>EPS):
        return (True, p1 + t*p2, jnp.cross(edge1,edge2))

    return (False, None, 0)

def render(edges, verts, grid_res):
    
    p1 = jnp.array([-2,0.2,0])
    light = jnp.array([-2,0,0.5])
    light = light/jnp.linalg.norm(light)
    
    vals = jnp.linspace(-1,1,grid_res)
    X,Y = jnp.meshgrid(vals, vals)
    rays = jnp.stack([jnp.ones((grid_res**2)),X.reshape(-1), Y.reshape(-1)]).T
    
    edges = device_put(edges)
    verts = device_put(verts)
    
    def get_ray(p2):
        
        dists = []
        for v in verts:
            dists.append(moeller_trumbore(p1,p2,edges[v[0]],edges[v[1]],edges[v[2]]))
                          
        normals = [x[2] for x in dists if x[0]]
        dists = [x[1] for x in dists if x[0]]
        
        if len(dists)>0:
            
            dists = jnp.sum((jnp.array(dists)-p1)**2,-1)
            dmin = jnp.argmin(dists)
            
            val = max(0,jnp.dot(normals[dmin]/jnp.linalg.norm(normals[dmin]), light))
            
            return val
        
        else:
            return -1
        
    result = jax.vmap(get_ray)(rays)
                
    return np.array(result)

I am basically just trying to use to jax to parallelize a very simple implementation of ray-tracing. I do not require automatic-differentiation (at least at the moment).

I understand that if jnp.abs(a) < EPS is causing the problem, because a is a function of p2 and vmap introduces abstract tracers for its arguments and p2 is the argument of the vmap. However, I don't know how I would fix this issue.

Any pointers greatly appreciated.

The issues I was having were caused by not using jax.lax.cond. Below you can find the fixed version.

The code below works fine but seems to be mainly executed on the CPU.

@jax.jit
def moeller_trumbore(p1, p2, q1, q2, q3):

    EPS = 1E-6
    
    vertex0 = q1
    vertex1 = q2
    vertex2 = q3
    
    edge1 = vertex1 - vertex0
    edge2 = vertex2 - vertex0
    
    h = jnp.cross(p2,edge2)
    a = jnp.dot(edge1,h)
    
    def compute_t(t):
        
        p1, p2, f, e1, e2, q = t
        t = f*jnp.dot(e2,q)
        
        return jax.lax.cond(t>1E-6, true_fun = lambda x: (True, p1 + t*p2, jnp.cross(e1,e2)),
                            false_fun = lambda x: (False,p1,p1),
                            operand=None)
    
    
    def compute_v(t):
        s, e1, e2, p1, p2, u, f = t
        
        q = jnp.cross(s,e1)
        v = f*jnp.dot(p2,q)
        
        #(v<0) or (u + v > 1.0)
        return jax.lax.cond((v<0) | (u + v > 1.0), true_fun = lambda x: (False,p1,p1), 
                            false_fun = compute_t,
                            operand = (p1,p2, f, e1, e2, q))
    
    
    def compute_u(t):
        
        a, p1, p2, v0, e1, e2, h = t
        
        f = 1./a
        s = p1 - v0
        u = f*jnp.dot(s,h)
        
        #(u<0) or (u>1.0)
        return jax.lax.cond((u<0) | (u>1.), true_fun = lambda _: (False,p1,p1),
                            false_fun = compute_v,
                            operand = (s, e1, e2, p1, p2, u, f))
    
        
    return jax.lax.cond(jnp.abs(a) < EPS, true_fun = lambda _: (False,p1,p1),
                        false_fun=compute_u,
                        operand = (a, p1, p2, vertex0, edge1, edge2, h))
    

def get_ray(p2, p1, verts, edges, light):
    
    min_dist = jnp.inf
    cval = -1.
    
    for v in verts:
        tf, p, normal = moeller_trumbore(p1,p2,edges[v[0]],edges[v[1]],edges[v[2]])
        
        d = jnp.sum((p-p1)**2)
        
        val = jnp.dot(normal/jnp.linalg.norm(normal), light)
        val = jnp.max(jnp.array([val,0]))
        
        min_dist, cval = jax.lax.cond(tf & (d<min_dist), true_fun=lambda _: (d, val),
                                      false_fun=lambda _: (min_dist, cval),
                                      operand = None)
        
        
    return cval

def render(edges, verts, grid_res):
    tik = time.time()
    
    p1 = jnp.array([-3.,0.,1.5])
    light = jnp.array([-2,1,1.5])
    light = light/jnp.linalg.norm(light)
    
    X,Y = jnp.meshgrid(jnp.linspace(-1,1,grid_res), jnp.linspace(-1.5,0.5,grid_res))
    rays = jnp.stack([jnp.ones((grid_res**2)), X.reshape(-1), Y.reshape(-1)]).T
    
    edges = device_put(edges)
    verts = device_put(verts)
    
    print("ray loc", rays.device_buffer.device())
    print("edg loc", edges.device_buffer.device())
    print("verts loc", verts[0][0].device_buffer.device())

    c = np.array([46., 140., 34.])
    ray_fun = lambda x: get_ray(x, p1, verts, edges, light)
    
    result = jax.vmap(ray_fun)(rays)
    print("res loc", result.device_buffer.device())
    
    result = np.expand_dims(np.array(result),-1)*c
    result[result<0] = 255.
    print("Time", time.time()-tik)
    return np.flipud(np.reshape(result,(grid_res,grid_res,3))).astype(np.uint8)

[jakevdp]

Any place that you use value-dependent control flow will cause this issue; I see several places in your snippet that would be problematic under JIT compilation. This doc page: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-Control-Flow has some discussion about why this is, and some ideas of what you can do to work around it.

[helange23]

Thanks! I got it to work now. It was very confusing that jax.lax.cond doesn't accept compound predicates, i.e. you need two cond-statements if you want to express something like if (u<0) or (u>1.0)

It's incredibly fast though!

[mattjj]

I think that's because the Python short-circuit boolean operators like or can't be overloaded, but you can instead write it using the | operator like (u<0) | (u>1.0).

If you write microbenchmarks, don't forget to account for async dispatch.

[helange23]

That makes sense. Thanks a lot.

One last question. When I run the (fixed) code above, the GPU utility is only about 1% and CPU utility is 100%.
I checked the placement of the variables with rays.device_buffer.device() and they are all on the GPU. The results are on the GPU as well.

How can this be if computation follows the placement of data? Am I missing something?

[jakevdp]

Could you edit the question above and add the fixed code for which you're seeing this performance issue? That would make your followup question easier to answer.