Is jit(f).lower() supposed to be callable inside jit?

[SJSeckmeyer]

I am not sure whether this behavior is a bug, whether you are not supposed to call lower inside a traced context or whether I overlook something else.

When you call jax.jit(fun).lower() on a function fun which gets passed one of its arguments via tree_util.Partial (or functools.partial for that matter) and this passed argument is a dynamic tracer, the call to lower crashes with a leaked tracer error:
jax._src.traceback_util.UnfilteredStackTrace: jax._src.errors.UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type int32[] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.

This error does not appear when one either:

1. calls the jitted function directly instead of calling lower.
2. passes the tracer as a normal argument instead of passing it via tree_util.Partial.

As far as I understand each tracer is annotated with a jit-level in order to detect escaping tracers. I suspect that in this case the tracer-level does not get incremented when b is passed down into f1. But maybe I am missing something else here.

Code to reproduce:

import jax
jax.config.update('jax_check_tracer_leaks', True)

def f1(a, b):
    a = a+b

@jax.jit
def f2(a, b):
    # works
    jax.jit(f1)(a,b)
    # works
    jax.jit(f1).lower(a,b)

    f1_partial = jax.tree_util.Partial(f1, b=b)
    # works
    jax.jit(f1_partial)(a)
    # crashes with error about leaked tracer.
    jax.jit(f1_partial).lower(a)

f2(0,1)

In case you wonder why I want to do this: I currently use lower to get the out_tree of a function. I use then this out_tree to pass in a matching input into scan. That way I can move computations out of the hot function that is called with scan.
This is like a custom "weak type upcast" which happens inside scan.

[mattjj]

Thanks for the question!

Producing such a bad error message is certainly a bug, but I believe the answer is no, .lower should only be used at "top-level" (i.e. not underneath jax.jit or any other JAX transformation or staging mechanism).

@froystig can you correct me if I'm wrong, and also think about how we might improve the error message?

[mattjj]

I currently use lower to get the out_tree of a function. I use then this out_tree to pass in a matching input into scan.

I believe the canonical way to do this is to use jax.eval_shape. You can pair it with jax.tree_util.tree_structure if you just want the output tree structure.

[froystig]

Yeah, we can and should improve the error message.

.lower should only be used at "top-level" (i.e. not underneath jax.jit or any other JAX transformation or staging mechanism).

It's the AOT-compiled function given by .lower(...).compile() that should only be called at "top level." I'm not sure that this constraint must also necessarily be placed on lowering (.lower(...)), although I don't see much use in allowing it either. For example, the following is a bit ridiculous but possible today:

>>> import jax
>>> import jax.numpy as jnp
>>> @jax.jit
... def f(x):
...   return x + jax.jit(lambda x: x).lower(x + jnp.ones(3)).in_avals[0][0].ndim
... 
>>> f(3)
Array(4, dtype=int32, weak_type=True)

The error in the OP is due to something more specific: lowering interacting with a closed-over tracer (which is readily available because of creating that closure within a transformation).

[SJSeckmeyer]

I believe the canonical way to do this is to use jax.eval_shape. You can pair it with jax.tree_util.tree_structure if you just want the output tree structure.

Thank you for the explanation and recommendation! jax.eval_shape is exactly what I was looking for.