jax
6168f5cd - [PjRt-IFRT] Tighten `PjRtExecutable::Create()` to take an MLIR module

Commit
16 days ago
[PjRt-IFRT] Tighten `PjRtExecutable::Create()` to take an MLIR module `xla::ifrt::PjRtExecutable::Create()` was taking already constructed `xla::PjRtExecutable`. This prevents `xla::ifrt::PjRtExecutable` from extracting useful information such as whether a default layout is used for outputs that is not available in the already constructed `xla::PjRtExecutable`. By taking a module and performing compilation inside, the `Create()` has access to the original MLIR module information. This is a prerequisite for improving the serialization of `xla::ifrt::PjRtLoadedExecutable` and `xla::ifrt::PjRtExecutable` to use `xla::ifrt::SerializedXlaExecutableMetadata` header. Ultimately, the richer executable metadata information will be used for unifying semantics for default layout handling across IFRT runtimes. PiperOrigin-RevId: 840491451
Author
Parents
Loading