Is it possible to access the AOT-compiled function its HLO structure, not only text?

[JTaets]

Hi all,

From the AOT-compilation docs, we can read the HLO as a string.
But it doesn't seem that we can get the "graph structure", i.e. list of operands and inputs like we can with a jaxpr?
I found some things in jaxlib, but i'm not sure.

This together with this doesn't seem to work anymore (and is installing tensorflow really required to be able to do this?)

Kind regards

[hawkinsp]

These days, it's probably best to use the stablehlo MLIR object, which you can get like this:

jax.jit(lambda x: x+2).lower(7).compiler_ir()

This returns an MLIR ir.Module object, which you can inspect and iterate over. Hope that helps!

[JTaets]

I can access the info I want from that, thanks for the response!

If I understand correctly, stablehlo is at a higher less-optimized level than hlo.
So is it possible to do more optimizations on stablehlo and still return accessible results?