Is there a way to know if I have accidently used an impure function with @jit?

[patel-zeel]

Hi, I am very new to jax and I encountered one issue working with @jit functions. As per documentation, I saw that jax operations work best on pure functions. However is there a way for users to know if they have accidently used @jit over an impure function? Below is the minimum working example of my problem:

from jax import jit

@jit
def function(x1, x2):
  return x1 * x2 if op == '*' else x1 + x2

op = '+'
print(function(2, 3))
op = '*'
print(function(2, 3))

Output

5
5

In the output, I do not see any warning or error showing that something went wrong here.

[YouJiacheng]

Maybe impossible.
JAX JIT use tracing instead of parsing AST, thus it cannot distinguish global variables from local variables.
But you can use some python code static analyzer to do this stuff.