How to generate a sub-jaxpr?

[YouJiacheng]

Let's say we have a function f, which call g. I'd like to get jaxpr of f, which call jaxpr of g instead of tracing through it. How can I achieve it?

In torch.fx, I can simply override is_leaf_module method of my custom Tracer, or use torch.fx.wrap.

But it seems harder to customize tracer in JAX(actually because the documentation about tracer is nearly blank), and the counterpart/alternative of torch.fx.wrap in JAX currently is registering user function as primitive.
However, to the best of my knowledge, JAX doesn't provide public API or documentation to help user, especially doesn't provide tools for automatically generating jvp, batching etc. rules for "primitive" implemented by other primitives.

In addition, though the usage jax.core.jaxpr_as_fun seems trivial, there isn't any documentation about it, while it is an important API for users who want to extend JAX's function transformation.

Further question: Is JAX Team planing to improve the (user-oriented) extensibility of the core of JAX? i.e. make Tracer can be extended by user more easily, provide high-level APIs/tools like torch.fx.subgraph_rewriter for implementing custom function transformation, etc. The APIs/tools and documentation

[sharadmv]

This is not a specific answer to the question of generating a "sub-jaxpr" but I'm interested in hearing more about your use case to better understand the need for subjaxprs.

Is JAX Team planing to improve the (user-oriented) extensibility of the core of JAX?

We are currently exploring this in this branch (take a look at example usage here). However, this is highly experimental and we'd like to understand use-cases more before we are to land this as a feature.

[YouJiacheng]

My motivation/use case was described in #10974 , I want to do function-level incremental compilation. i.e. I only change g in development, but JAX always compile the whole f.
I think if we have function-level incremental compilation, JAX's painfully slow compilation during developing can be remedied to a large extent.("sub-jaxpr", "sub XLA module" and persistent jit cache are needed)
Another use case was pointed out by patrick-kidger in that issue as well: g can be called many times, and cannot be expressed by scan(unlike the toy example provided in the comment).

[YouJiacheng]

I'd like to have an interface like torch.fx.wrap to make JAX don't trace through and inline g.