Fix segfault

PiperOrigin-RevId: 686821923
jaro-sevcik · Oct 17, 2024 · 9027fb3 · 9027fb3
1 parent 96d5542
commit 9027fb3
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -3868,6 +3868,10 @@ LogicalResult vector_multi_reduction_rule(RewriteContext &ctx, Operation &op,
                   src_layout.tileDataBounds(builder.getContext(), src_shape,
                                             src_idx, ctx.target_shape,
                                             allow_replicated);
+              if (data_bounds == nullptr) {
+                // Op error has already been emitted inside tileDataBounds().
+                return absl::UnknownError("Unable to obtain data bounds");
+              }
               // TODO(tlongeri): Maybe assemble/disassemble should take
               // TypedValue<VectorType> and we could save casts here and
               // elsewhere