Function is not jit-able when jnp.arange contain stop argument of jnp.maximum output

[riven314]

I have little clue on what is going wrong, so would like to put it here for advice:

I have a function (lets call it func1) which is jit-able (it is simplified to illustrate the issue):

import numpy as np
import jax
import jax.numpy as jnp

def func1(x):
    r = x.shape[0]
    D = jnp.zeros((r, r))
    for i in range(x.shape[0]):
        y = jnp.arange(i)
        D = jax.ops.index_update(D, y, i)
    return D

s1 = np.random.randn(10)
jit_func1 = jax.jit(func1)

# NO ERROR, SO FAR SO GOOD
_ = jit_func1(s1)

But once I made the change in jnp.arange, it is no longer jit-able:

def func2(x):
    r = x.shape[0]
    D = jnp.zeros((r, r))
    for i in range(x.shape[0]):
        # THIS IS THE CHANGE
        # replace stop argument by jnp.maximum output
        y = jnp.arange(jnp.maximum(0., i))
        D = jax.ops.index_update(D, y, i)
    return D

jit_func2 = jax.jit(func2)

# THIS WILL THROW AN ERROR
_ = jit_func2(s1)

The error is as follows:

UnfilteredStackTrace: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
It arose in jax.numpy.arange argument `stop`.
While tracing the function func2 at <ipython-input-6-84cc8f817b91>:10 for jit, this value became a tracer due to JAX operations on these lines:

  operation g:float32[] = convert_element_type[ new_dtype=float32
                                    weak_type=True ] f:int32[]
    from line <ipython-input-6-84cc8f817b91>:14 (func2)

  operation h:float32[] = max e:float32[] g:float32[]
    from line <ipython-input-6-84cc8f817b91>:14 (func2)

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

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:

ConcretizationTypeError                   Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py in arange(start, stop, step, dtype)
   3246   msg = "It arose in jax.numpy.arange argument `{}`.".format
   3247   if stop is None and step is None:
-> 3248     start = require(start, msg("stop"))
   3249     dtype = dtype or _dtype(start)
   3250     return lax.iota(dtype, np.ceil(start).astype(int)) # avoids materializing

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=0/1)>
It arose in jax.numpy.arange argument `stop`.
While tracing the function func2 at <ipython-input-6-84cc8f817b91>:10 for jit, this value became a tracer due to JAX operations on these lines:

  operation g:float32[] = convert_element_type[ new_dtype=float32
                                    weak_type=True ] f:int32[]
    from line <ipython-input-6-84cc8f817b91>:14 (func2)

  operation h:float32[] = max e:float32[] g:float32[]
    from line <ipython-input-6-84cc8f817b91>:14 (func2)

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

In both function, y has dynamic shape over loop. But I have no idea why the first one is jit-able but the second one is not. Could anyone shed me some lights as in what is going wrong for func2?

[jakevdp]

Hi - thanks for the question!

I think the definition of "dynamic" here is what is causing confusion. The issue of dynamic vs. static is not about whether a variable is changing during the program, but rather centers around what can be known by JAX at trace time. When you write a loop (like for i in range(5)), the value of i is known statically at trace-time.

However, the result of a jax operation (like jnp.maximum) is not evaluated at trace-time, even if its inputs are static. The output of the function is dynamic, and since JAX arrays must be statically shaped, you cannot create an array with a size that comes from the result of a jax function like jnp.maximum.

Does that answer your question?

[riven314]

Thanks for the prompt reply and pointing out to me the confusion!
It looks clearer to me now. Can I say that even though the shape of an array that is intermediately created changes over loop, it is still jit-able so long as its shape can be inferred by JAX in trace time?

[jakevdp]

Yes, I think that’s a good way to put it