Reverts 12e07c9

PiperOrigin-RevId: 761946361

Author: danielsuo

jax/_src/compiler.py

@@ -291,6 +291,19 @@ def backend_compile(
     executable_devices: xc.DeviceList,
     options: xc.CompileOptions,
     host_callbacks: Sequence[Any],
+) -> xc.LoadedExecutable:
+  return backend_compile_and_load(
+      backend, module, executable_devices, options, host_callbacks
+  )
+
+
+@profiler.annotate_function
+def backend_compile_and_load(
+    backend: xc.Client,
+    module: ir.Module,
+    executable_devices: xc.DeviceList,
+    options: xc.CompileOptions,
+    host_callbacks: Sequence[Any],
 ) -> xc.LoadedExecutable:
   sym_name = module.operation.attributes['sym_name']
   module_name = ir.StringAttr(sym_name).value
@@ -315,18 +328,35 @@ def backend_compile(
   try:
     # we use a separate function call to ensure that XLA compilation appears
     # separately in Python profiling results
-    if host_callbacks:
+    elif jaxlib_extension_version < 342 or isinstance(backend, xc.CompileOnlyPyClient):
+      if host_callbacks:
+        return backend.compile(
+            built_c,
+            executable_devices=executable_devices,  # type: ignore
+            compile_options=options,
+            host_callbacks=host_callbacks,
+        )
+      # Some backends don't have `host_callbacks` option yet
+      # TODO(sharadmv): remove this fallback when all backends allow `compile`
+      # to take in `host_callbacks`
       return backend.compile(
+          built_c, executable_devices=executable_devices, compile_options=options)  # type: ignore
+    else:
+      if host_callbacks:
+        return backend.compile_and_load(
+            built_c,
+            executable_devices=executable_devices,
+            compile_options=options,
+            host_callbacks=host_callbacks,
+        )
+      # Some backends don't have `host_callbacks` option yet
+      # TODO(sharadmv): remove this fallback when all backends allow `compile`
+      # to take in `host_callbacks`
+      return backend.compile_and_load(
           built_c,
-          executable_devices=executable_devices,  # type: ignore
+          executable_devices=executable_devices,
           compile_options=options,
-          host_callbacks=host_callbacks,
       )
-    # Some backends don't have `host_callbacks` option yet
-    # TODO(sharadmv): remove this fallback when all backends allow `compile`
-    # to take in `host_callbacks`
-    return backend.compile(
-        built_c, executable_devices=executable_devices, compile_options=options)  # type: ignore
   except xc.XlaRuntimeError as e:
     for error_handler in _XLA_RUNTIME_ERROR_HANDLERS:
       handler_result = error_handler(e)
@@ -391,7 +421,7 @@ def compile_or_get_cached(
   )
 
   if cache_key is None:
-    return backend_compile(
+    return backend_compile_and_load(
         backend, computation, executable_devices, compile_options,
         host_callbacks)
 
@@ -419,7 +449,7 @@ def compile_or_get_cached(
       config.share_binary_between_hosts.value
       and is_multi_process
       and distributed.global_state.client is not None
-      # Host callbacks are currently baked into the HLO module so we cant share
+      # Host callbacks are currently baked into the HLO module so we can't share
       # them.
       and len(host_callbacks) == 0
   ):
@@ -705,7 +735,7 @@ def _compile_and_write_cache(
     cache_key: str,
 ) -> xc.LoadedExecutable:
   start_time = time.monotonic()
-  executable = backend_compile(
+  executable = backend_compile_and_load(
       backend, computation, executable_devices, compile_options, host_callbacks
   )
   compile_time = time.monotonic() - start_time

jaxlib/_jax/__init__.pyi

@@ -551,6 +551,17 @@ class Client:
   ) -> PjRtLayout: ...
   def __getattr__(self, name: str) -> Any: ...
 
+
+class CompileOnlyPyClient(Client):
+  def compile(
+      self,
+      computation: str | bytes,
+      executable_devices: DeviceList | Sequence[Device],
+      compile_options: CompileOptions = ...,
+      host_callbacks: Sequence[Any] = ...,
+  ) -> LoadedExecutable: ...
+
+
 class CpuCollectives: ...
 
 def make_gloo_tcp_collectives(

jaxlib/xla_client.py

@@ -304,6 +304,7 @@ def computation_count():
 
 XlaComputation = _xla.XlaComputation
 Client = _xla.Client
+CompileOnlyPyClient = _xla.CompileOnlyPyClient
 Memory = _xla.Memory
 Array = _xla.Array
 ArrayImpl = _xla.ArrayImpl

jaxlib/xla_client.pyi

@@ -24,6 +24,7 @@ from jaxlib._jax import ArrayCopySemantics as ArrayCopySemantics
 from jaxlib._jax import ArrayImpl as ArrayImpl
 from jaxlib._jax import AutotuneCacheMode as AutotuneCacheMode
 from jaxlib._jax import Client as Client
+from jaxlib._jax import CompileOnlyPyClient as CompileOnlyPyClient
 from jaxlib._jax import CompileOptions as CompileOptions
 from jaxlib._jax import Device as Device
 from jaxlib._jax import DeviceAssignment as DeviceAssignment