Skip to content

Commit

Permalink
Include the sdy MLIR dialect in jaxlib.
Browse files Browse the repository at this point in the history
We're seeing test failures from tests assuming that this dialect exists. But given we plan to enable it at some point, we may as well just include it in the build.

The size impact is small (around 400K uncompressed).

PiperOrigin-RevId: 679592303
  • Loading branch information
hawkinsp authored and Google-ML-Automation committed Sep 27, 2024
1 parent 26632fd commit 6e1d850
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions jaxlib/tools/build_wheel.py
Original file line number Diff line number Diff line change
Expand Up @@ -351,6 +351,7 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu, skip_gpu_kernels):
f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsSparseTensor.{pyext}",
f"__main__/jaxlib/mlir/_mlir_libs/_mlirSparseTensorPasses.{pyext}",
f"__main__/jaxlib/mlir/_mlir_libs/_tpu_ext.{pyext}",
f"__main__/jaxlib/mlir/_mlir_libs/_sdy.{pyext}",
f"__main__/jaxlib/mlir/_mlir_libs/_stablehlo.{pyext}",
f"__main__/jaxlib/mlir/_mlir_libs/register_jax_dialects.{pyext}",
f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsGPU.{pyext}",
Expand Down

0 comments on commit 6e1d850

Please sign in to comment.