[jaxlib] Add CompileOnlyPyClient to xla_client.

We have users of CompileOnlyPyClient that use `backend.compile` as we eventually intend it (i.e., return `ExecutableRef`, possibly `PyExecutable` eventually, instead of `PyLoadedExectuable`).

PiperOrigin-RevId: 762440439

Author: danielsuo

jaxlib/_jax/__init__.pyi

@@ -553,6 +553,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 = 344
+_version = 345
 
 # An internal increasing version number for protecting jaxlib code against
 # ifrt changes.