Add support for deserializing xplanes to Jaxlib

jax/_src/lib/__init__.py

@@ -92,6 +92,7 @@ def _parse_version(v: str) -> tuple[int, ...]:
 from jaxlib._jax import pytree as pytree  # noqa: F401
 from jaxlib._jax import Device as Device  # noqa: F401
 from jaxlib import _profiler as _profiler  # noqa: F401
+from jaxlib import _profile_data as _profile_data  # noqa: F401
 
 import jaxlib.xla_client as xla_client  # noqa: F401
 

jax/_src/profiler.py

@@ -33,6 +33,11 @@
 
 from jax._src import xla_bridge
 from jax._src.lib import _profiler
+from jax._src.lib import _profile_data
+
+ProfileData = _profile_data.ProfileData
+ProfileEvent = _profile_data.ProfileEvent
+ProfilePlane = _profile_data.ProfilePlane
 
 _profiler_server: _profiler.ProfilerServer | None = None
 

jax/profiler.py

@@ -17,7 +17,10 @@
 from typing import Any
 
 from jax._src.profiler import (
+    ProfileData as ProfileData,
+    ProfileEvent as ProfileEvent,
     ProfileOptions as ProfileOptions,
+    ProfilePlane as ProfilePlane,
     StepTraceAnnotation as StepTraceAnnotation,
     TraceAnnotation as TraceAnnotation,
     annotate_function as annotate_function,
@@ -29,9 +32,3 @@
     stop_trace as stop_trace,
     trace as trace,
 )
-
-# this is a temporary shim to please pytype in the meantime before the migration
-# is complete for cl/760646494
-ProfileData: Any = None
-ProfileEvent: Any = None
-ProfilePlane: Any = None

jaxlib/BUILD

@@ -87,6 +87,7 @@ pytype_strict_library(
         "//jaxlib/mosaic/python:gpu_dialect",
         "//jaxlib/mosaic/python:tpu_dialect",
         "//jaxlib/triton",
+        "@xla//xla/python:_profile_data",
         "@xla//xla/python:_profiler",
     ],
 )
@@ -169,6 +170,7 @@ pywrap_library(
         "//jaxlib/mlir/_mlir_libs:_stablehlo",
         "//jaxlib/mlir/_mlir_libs:_tpu_ext",
         "//jaxlib/mlir/_mlir_libs:_triton_ext",
+        "@xla//xla/python:_profile_data",
         "@xla//xla/python:_profiler",
     ],
 )

jaxlib/tools/build_wheel.py

@@ -212,6 +212,7 @@ def prepare_wheel(wheel_sources_path: pathlib.Path, *, cpu, wheel_sources):
           f"{source_file_prefix}jaxlib/weakref_lru_cache.pyi",
           f"{source_file_prefix}jaxlib/_jax.{pyext}",
           f"{source_file_prefix}jaxlib/_profiler.{pyext}",
+          f"{source_file_prefix}jaxlib/_profile_data.{pyext}",
       ],
   )
   # This file is required by PEP-561. It marks jaxlib as package containing