How to jit function with conditional inside of for loop, and is it possible to chain functions before using jax.grad()?

[MyeloidSol]

Hello!
I've just recently started working with jax and was hoping for some help.

I'm having trouble figuring out how to jit this function:

# Calculate first principal direction
@partial(jit, static_argnums = (1,2,3))
def find_pd(mat, nfeat, max_iter, tol, key):
    # Split the PRNG key
    key = rand.split(key, 1)[0]

    # Initialize random vector
    cur_v = rand.uniform(key, shape = (nfeat,))

    # Find principal direction
    for _ in jnp.arange(0, max_iter):
        old_v = cur_v.copy()

        # Update current vector with gradient
        cur_v = cur_v + gradient(mat, cur_v)
        cur_v = cur_v / norm(cur_v)

        # Check if tolerance has been reached
        pred = norm(cur_v - old_v) < tol
        if pred: # <-- ERROR POPS UP HERE
            break
    
    # Return principal direction
    return( cur_v )

I get this error:

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
The error occurred while tracing the function find_pd at <ipython-input-38-abc964cc1074>:2 for jit. This value became a tracer due to JAX operations on these lines:

  operation a:f64[10] = pjit[
  name=_uniform
  jaxpr={ lambda ; b:key<fry>[] c:f64[] d:f64[]. let
      e:f64[] = convert_element_type[new_dtype=float64 weak_type=False] c
      f:f64[] = convert_element_type[new_dtype=float64 weak_type=False] d
      g:f64[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] e
      h:f64[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] f
      i:u64[10] = random_bits[bit_width=64 shape=(10,)] b
      j:u64[10] = shift_right_logical i 12
      k:u64[10] = or j 4607182418800017408
      l:f64[10] = bitcast_convert_type[new_dtype=float64] k
      m:f64[10] = sub l 1.0
      n:f64[1] = sub h g
      o:f64[10] = mul m n
      p:f64[10] = add o g
      q:f64[10] = max g p
    in (q,) }
] r s t
    from line <ipython-input-38-abc964cc1074>:8:12 (find_pd)

  operation a:f64[] = convert_element_type[new_dtype=float64 weak_type=False] b
    from line <ipython-input-38-abc964cc1074>:19:15 (find_pd)
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

I don't know how to rewrite the if-statement such that I can break out of the for-loop when tolerance is reached, I thought the tol variable was the issue but even when I get rid of it within the function and set it to some constant I still get an issue, then I tried using lax.cond() but I don't fully know how to call break within it, so I've ran into a dead-end here and am hoping for some help.

My second question is about this set of functions:

# Project a set of (column) vectors onto another vector
@jit
def orth_proj(mat, vec):
    return( jnp.dot(mat, vec) )

# Variance of projected (column) vectors
@jit
def op_var(mat, vec):
        return( orth_proj(mat, vec).var() )

# Gradient of projected variance w.r.t the input vector
gradient = jit(grad(op_var, argnums = 1))

Is there any way to do something like:

# Gradient of projected variance w.r.t the input vector
gradient = jit(grad(jnp.var(orth_proj), argnums = 1))

where instead of making the intermediate function op_var(), I can just "chain" var() after orth_proj and have the gradient still work properly? I just want to avoid needing the op_var function when all it's doing is just applying .var()

[patrick-kidger]

Question 1: if statements

For your first question: the if statement causes an error is because in JAX, we use Python to define the computation graph, not to run the computation graph. Here, a "computation graph" is simply some graph of operations that happen on JAX arrays. (And it is this graph which is JIT-compiled for speed.)

So having an if statement corresponds to having two different kinds of computation graph that we might want to define. Meanwhile a jax.lax.cond corresponds to actually doing branching within the graph itself.

(If you've any experience with metaprogramming in other languages, then that's basically what's going on here: we use Python at the metaprogramming level; the computation graph of JAX arrays is then the program itself. If you're not familiar with metaprogramming then don't worry about this analogy though!)

You might also like points 1 and 2 in this guide for a little more on this. (And I think the core JAX docs also have something about this somewhere.)

Question 2: breaking

Now, how to break a for loop in JAX? The answer is to write things as a while loop -- jax.lax.while_loop -- instead. Put the break condition into the condition function.

In particular note that we're using jax.lax.while_loop, and not a regular Python while loop. (See point 7 in the guide I link above for more on this.)

Question 3: avoiding the wrapper function

Unfortunately there's no way to avoid the wrapper function in general.

One thing you can do is to use a little lambda function directly where you need it. It's fairly common to see things like jax.grad(lambda x: f(x).var())(z), for example.

[MyeloidSol]

Thank you for the explanation and the link!

I tried to use lax.while_loop instead but for some reason I'm getting the same type of error:

# Calculate first principal direction
@partial(jit, static_argnums = (1,2,3))
def find_pd(mat, nfeat, max_iter, tol, key):
    # Split the PRNG key
    key = rand.split(key, 1)[0]

    # Initialize random vector
    cur_v = rand.uniform(key, shape = (nfeat,))

    # Find principal direction
    def cond_fun(state):
        # Unpack current iteration's state
        cur_v, diff, iter = state

        # Check tolerance and iteration
        return( diff < tol or iter > max_iter ) # <-- ERROR POPS UP HERE
    def body_fun(state):
        # Unpack last iteration
        cur_v, diff, iter = state
        iter += 1

        # Keep track of last iteration's vector
        old_v = cur_v.copy()

        # Update current vector with gradient
        cur_v = cur_v + gradient(mat, cur_v)
        cur_v = cur_v / norm(cur_v)

        # Difference to check if tolerance has been reached
        diff = norm(cur_v - old_v).item()

        return(cur_v, diff, iter)

    cur_v = lax.while_loop(cond_fun, body_fun, (cur_v, tol, 0))[0]

    # Return principal direction
    return( cur_v )

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..

I'm not sure why the issue persists I thought that stating tol is static will ensure that it's okay to use it like this, but now I'm unsure

[patrick-kidger]

You want (diff < tol) | (iter > max_iter).

or is a keyword in Python that operates on True and False.

| is an operator that can be overriden for custom types (in this case, JAX arrays).

[MyeloidSol]

Thank you so much! Doing this and getting rid of .item() in diff = norm(cur_v - old_v).item() makes it work perfectly!