Some questions about the workflow and the IR things of JAX

[Dong-Jiahuan]

Hello jax experts. I have some questions about the workflow and the IR things of JAX to ask:

1.When I JIT compile a JAX function, how does this function finally run on the device?
Am I in the right way to think as follows:
The function first be converted into jaxpr and then hlo and then mhlo. The mhlo representation is optimized in MLIR system (like operation fusion) and then converted into hlo again. The optimized hlo is sent to XLA and doesn't need be optimized by XLA (even though XLA can do this work) because MLIR had done this work. Or maybe the optimized hlo still need be optimized by XLA in some aspects of which MLIR didn't cover. Anyway, hlo is sent to XLA backend and XLA backend perfoems target-dependent optimization on the HLO-level and converts hlo into LLVM IR etc. All in all, JAX compile workflow is not totally in XLA because of the join of MLIR.

2.According to the change log (jax 0.2.27) , I can get the IR file by setting "JAX_DUMP_IR_TO" environment variable and the file I get is represented in mhlo. I want to know the mhlo I got is optimized or not? Can I get other form of IR file, like hlo ? If I can, how can I know it's optimized or not?

3.I also know using the api "jax.jit(your_function).lower(function_arguments...).as_text(dialect='hlo/mhlo')" can get the hlo/mhlo representation of a function. If I set dialect=mhlo, I think I'll get almost the same content with the IR file I got by setting "JAX_DUMP_TO_IR". (The only difference is that the IR file I got by setting environment variable has the information of "loc" but the string I got by using api doesn't have "loc" information) Am I thinking in the right way?

Thank you for your time and I'm looking forward to your reply!