How can I obtain the computation graph or MLIR file generated by JIT optimization?

[knightXun]

For code like : output2, _ = jax.jit(model.apply)(variables, input) , I want to obtain the XLA computation graph or mlir code to analyse the middle optimization. How can I do it?

[jakevdp]

You can output the compiled HLO using the Ahead of time lowering and compilation tools. It looks something like this:

print(jax.jit(model.apply).lower(variables, input).compile().as_text()

[knightXun]

how about the uncompiled HLO?

[jakevdp]

Remove .compile() to print the uncompiled HLO text

[knightXun]

How can I print out the xla optimization steps it has taken and specify that it only performs specific xla optimization steps?

[jakevdp]

I don't know of any way to access that information from the JAX python layer.

[knightXun]

Does JAX directly send the source code to XLA for translation, or is the code already translated into HLO and then handed over to XLA for optimization? If it is translated into HLO first, how can I obtain this HLO?

[jakevdp]

JAX translates the Python to HLO before sending it to XLA for optimization. You can obtain this HLO the way I mentioned above, removing .compile() from the original snippet:

print(jax.jit(model.apply).lower(variables, input).as_text()