Check whether JAX is currently tracing

[cool-RR]

I'm writing a function that should never be run with JIT. I want to add an assert line at the end that asserts that JAX is not tracing this function. How do I check for this?

[jakevdp]

JAX deliberately does not provide any API for this. Code with logic that branches on trace state is an anti-pattern and should be avoided: for example, it might cause autodiff to silently return wrong results if you use different codepaths when tracing and when not.

If your goal is to fail if a function is being traced, one way to do so would be to write code that fails under tracing; e.g. jnp.unique(0) would raise a ConcretizationTypeError under tracing. That may not be a very satisfying answer though.

[cool-RR]

I am actually feeling quite satisfied by this answer. Thank you!