[Pallas] Better error message for grid/dimension semantics mismatch.

justinjfu · Google-ML-Automation · commit e7545eca58d4 · 2025-11-03T08:52:25.000-08:00
PiperOrigin-RevId: 827516689
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
@@ -472,7 +472,8 @@ def __init__(
     )
     if len(user_grid) != len(dimension_semantics):
       raise ValueError(
-          "Must have dimension semantics for each dimension of the grid."
+          "Length of grid does not match length of dimension semantics."
+          f" len(grid)={len(user_grid)}, {len(dimension_semantics)=}"
       )
     assert len(self.vmapped_dims) + len(dimension_semantics) == len(
         self.grid

Original file line number	Diff line number	Diff line change
`@@ -472,7 +472,8 @@ def __init__(`
`472`	`472`	`)`
`473`	`473`	`if len(user_grid) != len(dimension_semantics):`
`474`	`474`	`raise ValueError(`
`475`		`- "Must have dimension semantics for each dimension of the grid."`
	`475`	`+ "Length of grid does not match length of dimension semantics."`
	`476`	`+ f" len(grid)={len(user_grid)}, {len(dimension_semantics)=}"`
`476`	`477`	`)`
`477`	`478`	`assert len(self.vmapped_dims) + len(dimension_semantics) == len(`
`478`	`479`	`self.grid`