Consider raising when a jitted function's output is not PyTree-like

Just as https://github.com/google/jax/issues/2813 argues that inputs to jit should always be hashable since that prevents unnecessary recompilation, another common error with jit is when its output is not PyTree-like (either a PyTree or registered using register_pytree_node). This has caught me a few times with tracers leaking out of functions. It makes me apprehensive about decorating functions with jit. Is there any use case for non-PyTree-like objects being returned? If not, I propose checking that the return value is a PyTree and raising if not.

[jekbradbury]

This is interesting. If I’m following correctly, we already do the same check for function arguments (we throw a “not a JaxType” error on unknown PyTree leaves). It seems like it would be sensible to also error on unknown PyTree leaves in return values, but there also might be valid uses for this (basically returning “compile-time” Python data).

[gnecula]

I think I understand. The following code

class MyClass:
  def __init__(self, x):
     self.x = x

jax.jit(lambda x: MyClass(x + 2))(0)

will return an instance of MyClass with the tracer inside. No code is jitted, because JAX does not see the returned Tracers. This is a confusing.

I think that James is thinking about an example like:

jax.jit(lambda x: x + 5, "hello")(0)

Maybe for such examples we can use a has_aux parameter to the jit to declare the the function being jitted as an intentional auxiliary result (a la grad)