[Mosaic] Fix lowering for _dot_general_lowering_rule to match the n…

…ew `vector.MultiDimReductionOp` signature. PiperOrigin-RevId: 662933072
jax-ml · Aug 14, 2024 · df2e9c3 · df2e9c3
1 parent b0a144a
commit df2e9c3
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/jax/_src/pallas/mosaic/lowering.py b/jax/_src/pallas/mosaic/lowering.py
@@ -1460,9 +1460,7 @@ def _dot_general_lowering_rule(
         ir.Attribute.parse("#vector.kind<add>"),
         arith.MulFOp(x, y),
         acc,
-        ir.ArrayAttr.get(
-            [ir.IntegerAttr.get(ir.IntegerType.get_signless(64), 1)]
-        ),
+        [1]
     )
     return vector.ShapeCastOp(out_type, red).result