fix

Author: wsmoses

patches/jax.patch

@@ -1,19 +1,21 @@
 diff --git a/jaxlib/mosaic/BUILD b/jaxlib/mosaic/BUILD
-index 088411939..598118390 100644
+index 088411939..ee6d632b6 100644
 --- a/jaxlib/mosaic/BUILD
 +++ b/jaxlib/mosaic/BUILD
-@@ -54,7 +54,15 @@ cc_library(
-         "dialect/tpu/tpu_dialect.h",
-         "dialect/tpu/util.h",
-         "dialect/tpu/vreg_util.h",
+@@ -48,13 +48,11 @@ cc_library(
+         "dialect/tpu/util.cc",
+         "dialect/tpu/vreg_util.cc",
+     ],
+-    hdrs = [
+-        "dialect/tpu/array_util.h",
+-        "dialect/tpu/layout.h",
+-        "dialect/tpu/tpu_dialect.h",
+-        "dialect/tpu/util.h",
+-        "dialect/tpu/vreg_util.h",
 -    ],
-+    ] + glob(
++    hdrs = glob(
 +        [
-+            "dialect/tpu/transforms/*.h",
-+        ],
-+        exclude = [
-+            "dialect/tpu/transforms/serde.h",
-+            "dialect/tpu/transforms/linalg_vectorization.h",
++            "dialect/tpu/**/*.h",
 +        ])
 +	,
      # compatible with libtpu