From d40d01819e21c7f8c1190cb48542d58e03320ca2 Mon Sep 17 00:00:00 2001
From: Matthias Kramm <kramm@google.com>
Date: Thu, 2 Oct 2025 13:24:33 -0700
Subject: [PATCH] Use less memory during MLIR compilation.

PiperOrigin-RevId: 814360386
---
 jaxlib/py_client.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/jaxlib/py_client.cc b/jaxlib/py_client.cc
index bd3c3eb8ebd5..7e87de889648 100644
--- a/jaxlib/py_client.cc
+++ b/jaxlib/py_client.cc
@@ -515,8 +515,9 @@ PyClient::CompileAndLoad(nb_class_ptr<PyClient> client, mlir::ModuleOp module,
       TF_RETURN_IF_ERROR(xla::ExportShardyForGSPMD(module));
     }
   }
+  options.allow_in_place_mlir_modification = true;  // we just cloned the module
   return CompileAndLoadIfrtProgram(
-      client, std::make_unique<xla::ifrt::HloProgram>(module),
+      client, std::make_unique<xla::ifrt::HloProgram>(std::move(module)),
       MakeIfrtCompileOptions(std::move(options), std::move(executable_devices),
                              std::move(host_callbacks)));
 }