How Jax lower jaxpr control flow into XLA?

[adamantboy]

how jaxpr_subcomp deal with control flow primitives, when these eqns are lowered into XLA HLO? For example, cond, while, scan.

[sharadmv]

Control flow primitives are lowered just like any other primitive in the sense that at lowering time, we are provided with 1) the MHLO arguments to the primitives and 2) the parameters of the primitive which include the jaxpr body (multiple jaxprs in the case of while/cond).

At a high level, we recursively call mlir.jaxpr_subcomp on these body jaxprs (see here for an example) while inside MLIR control flow blocks.

Does that answer your question or did you want some more specific answers for each primitive?

[adamantboy]

Thank you for your reply. One more question, what is the procedure of translating JAX into HLO? Is it just translated primitive by primitive?Is there any limitation during the lowering?It is important because JAX is a dynamic programming language while HLO is an IR of static programming language.

[sharadmv]

First, Python code that calls JAX primitives is converted into the Jaxpr format (see jax.make_jaxpr) via tracing, which is done primitive-by-primitive.Since we pass in values with known shapes/dtypes into the Python functions while tracing, the Jaxpr will have statically known shapes/dtypes as well. The jaxpr is thus specialized to the input shapes/dtypes to the function we are lowering to HLO.

Once we have a jaxpr, we now have a language that is more compatible with HLO and we don't need to worry about unknown shapes/dtypes. The inherent limitations are that we can't stage out arbitrary Python code, only Python code representable in jaxprs.