Why does jit + if statement work here?

[HHalva]

Going through some Jax source code, I see that a lot of functions use if statements instead of lax.cond (below example from one of the linear solvers). I am able to use these functions in jitted code and code runs fine -- I thought that if statements shouldn't work under jit? Or does it only work with is None statements as the tracing can be done here between the branches? I would like to understand this bit better.

def _isolve(_isolve_solve, A, b, x0=None, *, tol=1e-5, atol=0.0, maxiter=None, M=None, check_symmetric=False, has_info=False):
  if x0 is None:
    x0 = tree_map(jnp.zeros_like, b)

  b, x0 = device_put((b, x0)) 

...

[jakevdp]

jit is compatible with if statements whose truth or falsehood is decidable at trace time. x is None fits this criterion, because it concerns the identity of the Python variable.

Other statements that can be used within if statements are static attributes of the an array, such as shape, size, or dtype.

What will not work within JIT is attempting to branch based on dynamic quantites, such as values within arrays.

There's some discussion of these topics in How To Think In JAX.

[HHalva]

OK that makes sense -- have somehow not seen that linked site before, it's very useful. Thanks for a quick reply!