Fix jaxlib wheel build after removal of mosaic python files.

PiperOrigin-RevId: 597536465
google · Jan 11, 2024 · 858fd52 · 858fd52
1 parent 57506b5
commit 858fd52
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/jaxlib/BUILD b/jaxlib/BUILD
@@ -55,16 +55,20 @@ py_library_providing_imports_info(
         "//jaxlib/cpu:_ducc_fft",
         "//jaxlib/cpu:_lapack",
         "//jaxlib/mlir",
+        "//jaxlib/mlir:arithmetic_dialect",
         "//jaxlib/mlir:builtin_dialect",
         "//jaxlib/mlir:chlo_dialect",
         "//jaxlib/mlir:func_dialect",
         "//jaxlib/mlir:ir",
+        "//jaxlib/mlir:math_dialect",
         "//jaxlib/mlir:memref_dialect",
         "//jaxlib/mlir:mhlo_dialect",
         "//jaxlib/mlir:pass_manager",
         "//jaxlib/mlir:scf_dialect",
         "//jaxlib/mlir:sparse_tensor_dialect",
         "//jaxlib/mlir:stablehlo_dialect",
+        "//jaxlib/mlir:vector_dialect",
+        "//jaxlib/mosaic",
     ],
 )
 

diff --git a/jaxlib/tools/build_wheel.py b/jaxlib/tools/build_wheel.py
@@ -252,8 +252,6 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu, include_gpu_plugin_extensi
   copy_runfiles(
       dst_dir=mosaic_python_dir,
       src_files=[
-          "__main__/jaxlib/mosaic/python/apply_vector_layout.py",
-          "__main__/jaxlib/mosaic/python/infer_memref_layout.py",
           "__main__/jaxlib/mosaic/python/layout_defs.py",
           "__main__/jaxlib/mosaic/python/tpu.py",
       ],