jaxlib: Add ifrt_proxy.pyi to build_wheel.py.

jax authors · jax authors · commit df9cefabc14e · 2024-03-19T13:27:39.000-07:00
PiperOrigin-RevId: 617275734
diff --git a/jaxlib/tools/build_wheel.py b/jaxlib/tools/build_wheel.py
@@ -100,6 +100,7 @@ def patch_copy_mlir_import(src_file, dst_dir):
 
 _XLA_EXTENSION_STUBS = [
     "__init__.pyi",
+    "ifrt_proxy.pyi",
     "jax_jit.pyi",
     "ops.pyi",
     "outfeed_receiver.pyi",
