Rename backend.compile to backend.compile_and_load.

Part of a larger refactor. Today, `compile` returns a loaded executable i.e., fuses the compile and load functions. Eventually, `compile` should return an unloaded executable and `load` should return a loaded exectuable; the default jit path will still return a loaded executable.

PiperOrigin-RevId: 761951858

Author: danielsuo

jax/_src/compiler.py

@@ -34,7 +34,9 @@
 from jax._src import profiler
 from jax._src import traceback_util
 from jax._src.interpreters import mlir
+from jax._src.lib import jaxlib_extension_version
 from jax._src.lib import xla_client as xc
+from jax._src.lib import _jax
 from jax._src.lib.mlir import ir
 import numpy as np
 
@@ -291,6 +293,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 +330,40 @@ def backend_compile(
   try:
     # we use a separate function call to ensure that XLA compilation appears
     # separately in Python profiling results
-    if host_callbacks:
+    # TODO(dsuo): Simplify this logic once backend_compile actually returns an
+    # unloaded executable.
+    if jaxlib_extension_version < 343 or (
+        jaxlib_extension_version >= 343
+        and isinstance(backend, _jax.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 +428,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 +456,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 +742,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

@@ -43,7 +43,7 @@
 
 # Just an internal arbitrary increasing number to help with backward-compatible
 # changes. In JAX, reference this via jax._src.lib.jaxlib_extension_version.
-_version = 342
+_version = 343
 
 # An internal increasing version number for protecting jaxlib code against
 # ifrt changes.