Constant tracer

[ArtemBasyrov]

Hi,

I'm getting an error:

/jax/_src/pjit.py:214, in _python_pjit_helper(fun, jit_info, *args, **kwargs)
[212] (/jax/_src/pjit.py:212) # Run canonicalization again to figure out which arg failed.
    [213](/jax/_src/pjit.py:213) if p.params['jaxpr'].consts:
--> [214](/jax/_src/pjit.py:214)   raise TypeError(e.args[0]) from e
    [215](/jax/_src/pjit.py:215) else:
    [216](/jax/_src/pjit.py:216)   for arg, name, aval in zip(args_flat, arg_names, p.in_avals):

TypeError: Argument 'Traced<ShapedArray(float64[33], weak_type=True)>with<DynamicJaxprTrace>' of type <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'> is not a valid JAX type.

It seems like there are constants in my dynamic tracer, but I don't know how to trace where they are created. Is there a way to find where they are created?

[jakevdp]

Hi - thanks for the question! Depending on the context, this might be a bug, or it might be working as expected. If you could edit your question to add a minimal reproducible example of code that generates this error, we might be able to shed more light on it!