jnp.arange with Tracer

[sychang42]

Hi,
I am using jax and flax for a deep-learning problem.

During the training, I have a loss function which contains jnp.arange function with tracer as argument :

def epsilon_schedule(p, diameter, blur, scaling):
    eps_list = (
        [diameter ** p]
        + [
            jnp.exp(e)
            for e in jnp.arange(
                p * jnp.log(diameter), p * jnp.log(blur), p * jnp.log(scaling)
            )
        ]
        + [blur ** p]
    )
    return eps_list

where p, diameter, blur and scaling are essentially TracedArray.

However, this causes the ConcretizationTypeError as following :

> jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=0/3)>
It arose in jax.numpy.arange argument `start`. 

Would there be a way to circumvent this issue?

Thanks a lot!

[jakevdp]

JAX cannot act on dynamically-sized arrays, where by "dynamically sized" we mean arrays whose length depends on the value of traced inputs. Your code appears intended to do exactly this, and so it cannot be used in JAX. Typical workarounds involve either (1) marking problematic arguments as static, meaning they are unchanged from run to run, or (2) constructing fixed-length arrays with padding up to a maximum expected value. It's not clear to me whether either workaround is feasible in your case, because I'm not sure about the context in which the function is used.