What replaced as_serialized_hlo_module_proto() following the deprecation of jax.xla_computation?

[joaospinto]

Of course we can do something like f -> jit -> lower -> compile -> as_text(), but this is actually the equivalent of as_hlo_text(), not as_serialized_hlo_module_proto(), so cannot be easily de-serializable into a protobuf.

Also, even if I use this deprecated function call, when I load the protobuf using the PJRT C API, I get the following error:
INTERNAL: RET_CHECK failure (external/xla/xla/client/xla_computation.cc:27) proto_.has_host_program_shape()
This happens when I run this: https://github.com/joaospinto/call_jax_from_cpp/blob/main/cpp/main_hlo_to_serialized_c_api.cpp

You can see the code I'm using (everything is very much a minimal example, which is the purpose of the repo) here: https://github.com/joaospinto/call_jax_from_cpp/blob/main/call_jax_from_cpp/simple_jax_example.py.

Note that while the C++ PJRT API does handle as_hlo_text() outputs, the C PJRT API does not, so as_serialized_hlo_module_proto() is still relevant.

Also, what format is the output of f -> jit -> lower -> compile -> serialize in? Is it at all related to any of these protos?

[joaospinto]

CC @skye

[yashk2810]

Try: jax.jit(f).lower(*args).compiler_ir().as_serialized_hlo_module_proto()?

jax.jit(f).lower(*args).compiler_ir() will give you the same type back as xla_computation. So you can use whatever you were using before.

[joaospinto]

@yashk2810 The type of jax.xla_computation(f)(dummy) seems to be jaxlib.xla_extension.XlaComputation, while the type of jax.jit(f).lower(*args).compiler_ir() seems to be jaxlib.mlir._mlir_libs._mlir.ir.Module, which does not have a as_serialized_hlo_module_proto method:

AttributeError: 'jaxlib.mlir._mlir_libs._mlir.ir.Module' object has no attribute 'as_serialized_hlo_module_proto'

[yashk2810]

You also need to specify hlo i.e. jax.jit(f).lower(*args).compiler_ir('hlo')

[joaospinto]

@yashk2810 That works, thanks. I see that this commit made the error message explain what you just said. Would be awesome to have that in the docs, as well. That definitely answers the main question here.

Would be great if you could also clarify this part: in what format is the output of f -> jit -> lower -> compile -> serialize? E.g. is it a serialized proto (e.g. ExecutableAndOptionsProto) that can be directly consumed in the C++ APIs, or is it entirely unrelated?

[yashk2810]

there is no compile -> serialize step available on the AOT APIs.

Can you clarify via code, what you actually mean? There is a serialization option available on the Compiled object but that's experimental.

[joaospinto]

Clarifying: I meant to ask if the serialization in this minimal example is some internal JAX representation, or whether it corresponds to any of the XLA serialized proto formats. Looking at the code here, it seems to be completely independent of that.

[joaospinto]

@yashk2810 The last part of my initial question (the _has_host_program_shape error) is probably more of an XLA question; posted it here, feel free to take a look.