How to solve TracerIntegerConversionError?

[lanmao998]

I can run the following code without @jax.jit, but when I add @jax.jit, it failed. How can I solve this problem?

import jax.numpy as jnp
import jax

@jax.jit
def test(a):
    N = jnp.ceil(jnp.log(a)/jnp.log(12))
    N = N.astype('int')
    for i in range(N): 
        print(i)
test(123)

/usr/local/lib/python3.8/dist-packages/jax/_src/numpy/lax_numpy.py:4459: UserWarning: Explicitly requested dtype int requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.
  lax_internal._check_user_dtype_supported(dtype, "astype")
Traceback (most recent call last):
  File "test.py", line 12, in <module>
    test(123)
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/api.py", line 473, in cache_miss
    out_flat = xla.xla_call(
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 1765, in bind
    return call_bind(self, fun, *args, **params)
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 1781, in call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 678, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py", line 182, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/usr/local/lib/python3.8/dist-packages/jax/linear_util.py", line 285, in memoized_fun
    ans = call(fun, *args)
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py", line 230, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars, False,
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/dispatch.py", line 272, in lower_xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
  File "/usr/local/lib/python3.8/dist-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/jax/interpreters/partial_eval.py", line 1893, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File "/usr/local/lib/python3.8/dist-packages/jax/interpreters/partial_eval.py", line 1865, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File "/usr/local/lib/python3.8/dist-packages/jax/linear_util.py", line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "test.py", line 9, in test
    for i in range(N):
  File "/usr/local/lib/python3.8/dist-packages/jax/core.py", line 537, in __index__
    raise TracerIntegerConversionError(self)
jax._src.traceback_util.UnfilteredStackTrace: jax._src.errors.TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test.py", line 12, in <module>
    test(123)
  File "test.py", line 9, in test
    for i in range(N):
jax._src.errors.TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError

[YouJiacheng]

JAX doesn't support dynamic control flow. (specifically, Tracer doesn't support __index__ method used by range(n) or arr[i](for tuple and list) etc.)
There are 2 options:

1. Mark a as static argument. And use numpy instead of jnp for N computation, or use jax.ensure_compile_time_eval(not recommended).
2. Use jax.lax.while_loop instead of for i in range(N). see https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.while_loop.html

[lanmao998]

Thanks! It helps me a lot.

[tomsturges]

Would you be so kind as to explain the pros/cons of these two different methods?

[jakevdp]

Broadly: while_loop will lead to faster compilation, at the expense of slower runtime, with particular performance penalties on accelerators. Marking N as static will lead to slower compilation (growing roughly with N^2 as N gets large), but faster runtime. The best option would be to write your code in terms of vectorized operations rather than loops, but that is not always possible.