Issue running AOT-compiled code: proto.has_host_program_shape() No program shape found in the proto

[joaospinto]

Here's what I'm trying to do (each of these running in separate binaries):

1. Create an HLO file specifying some computation (in JAX, works OK).
2. Run AOT-compilation of this HLO-defined computation with XLA (works OK, but might not be serializing all that is needed).
3. Load the AOT-compiled computation and run it on a given example (crashes).

Here is the error I am observing:

2024-07-03 11:45:21.312474: E external/xla/xla/status_macros.cc:56] INTERNAL: RET_CHECK failure (external/xla/xla/hlo/ir/hlo_module.cc:551) proto.has_host_program_shape() No program shape found in the proto
*** Begin stack trace ***
        tsl::CurrentStackTrace()
        xla::status_macros::LogError(absl::lts_20230802::Status const&, char const*, int, int, bool)
        xla::status_macros::MakeError(char const*, int, absl::lts_20230802::StatusCode, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&, bool, int, bool)
        xla::status_macros::MakeErrorStream::Impl::GetStatus()
        xla::status_macros::MakeErrorStream::GetStatus() const
        xla::status_macros::MakeErrorStream::MakeErrorStreamWithOutput::operator absl::lts_20230802::StatusOr<std::__1::unique_ptr<xla::HloModule, std::__1::default_delete<xla::HloModule>>><std::__1::unique_ptr<xla::HloModule, std::__1::default_delete<xla::HloModule>>>()
        xla::HloModule::CreateFromProto(xla::HloModuleProto const&, xla::HloModuleConfig const&, bool)
        xla::HloModule::CreateFromProtoWithConfig(xla::HloModuleProtoWithConfig const&, bool)
        xla::cpu::(anonymous namespace)::CpuExecutableAotCompilationResult::FromString(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&)
        xla::cpu::CpuCompiler::LoadAotCompilationResult(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char>> const&)
        xla::TfrtCpuClient::DeserializeExecutable(std::__1::basic_string_view<char, std::__1::char_traits<char>>, std::__1::optional<xla::CompileOptions>)
        main
        start
*** End stack trace ***

libc++abi: terminating due to uncaught exception of type absl::lts_20230802::BadStatusOrAccess: Bad StatusOr access: INTERNAL: RET_CHECK failure (external/xla/xla/hlo/ir/hlo_module.cc:551) proto.has_host_program_shape() No program shape found in the proto

The code the steps listed above can be found respectively in:

1. https://github.com/joaospinto/call_jax_from_cpp/blob/main/call_jax_from_cpp/simple_jax_example.py (irrelevant for this question, but sharing for completeness)
2. https://github.com/joaospinto/call_jax_from_cpp/blob/main/cpp/main_hlo_to_serialized.cpp
3. https://github.com/joaospinto/call_jax_from_cpp/blob/main/cpp/main_aot.cpp
   Note that these are minimal examples.

This is related to the following discussion: google/jax#22184 (reply in thread)

Am I doing something wrong, or is something broken? Thanks!

[cheshire]

CC @ezhulenev