From 9748e2ab1a37a4587bdd648c458e1ec83c345d42 Mon Sep 17 00:00:00 2001 From: jax authors Date: Tue, 8 Oct 2024 15:06:29 -0700 Subject: [PATCH] [JAX] Fix error message for matmul operand shape check. PiperOrigin-RevId: 683778484 --- jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc index 2894b0797e7b..23409f231d84 100644 --- a/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc +++ b/jaxlib/mosaic/dialect/tpu/transforms/infer_vector_layout.cc @@ -856,7 +856,7 @@ class VectorLayoutInferer { auto shape = vty.getShape().take_back(2); if (shape[0] % major_multiple.value_or(tiling[0]) != 0 || shape[1] % minor_multiple.value_or(tiling[1]) != 0) { - op->emitOpError("Matmul operand") + op->emitOpError("Matmul operand ") << operand_name << " must have a shape divisible by (" << major_multiple.value_or(tiling[0]) << ", " << minor_multiple.value_or(tiling[1]) << "), but got: (" << shape[0]