condition function on jax

[lucasliunju]

Hi, I would like to use condition function (if... else...) to control data flow.

I have tried jax.numpy.where and jax.lax.cond

I find:
1. jax.numpy.numpy cannot be used to calculate gradients
2. jax.lax.cond can be used to calculate gradients, but it will calculate both the two conditions and that will cause extra time.
So I would like to ask whether there are some condition functions which just calculate the selected condition.

Thank you!

[jakevdp]

Hi - thanks for the question. I think we'll need more information to give you a useful answer. jnp.numpy.where is implemented using jax.lax.cond, so they should not differ in terms of their compatiblilty with autodiff.

Can you give a short, self-contained example of code that demonstrates the issue you're having?

[lucasliunju]

Thank you.

For example

def func(g_v, t):

    def d_1(v):
        g = jax.value_and_grad(loss_fn)(params, images, labels)
        g = jax.tree_map(lambda x: jax.lax.pmean(x, axis_name='batch'), g)
        return g
    
    def d_2(v):
        return v
    g = jax.lax.cond(t < 1000, d_1, d_2, g_v)

The condition function is used to determine whether we need to calculate the gradient, which can improve the efficiency compared with the method that calculate the gradient at each time step.

My question is whether jax.lax.cond can improve the efficiency and just calculate the gradient when the condition is true.
Or it calculates the results of two conditions and just use the result of true condition.
Thank you!

[jakevdp]

That makes sense, thanks. The answer is both yes and no!

If you call a function like this in normal eager mode, JAX will indeed fully compute both branches.

But if you JIT compile such a function, XLA can generally prune unnecessary computations from the computation graph, and only the necessary values will actually be computed.