[Mosaic] C++ apply-vector-layout: fix unnecessarily setting out_layou…

…t in matmul rule (which never existed in Python) PiperOrigin-RevId: 584568680
jax-ml · Nov 22, 2023 · 8457b02 · 8457b02
1 parent c9d8c4b
commit 8457b02
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc b/jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc
@@ -1181,7 +1181,6 @@ LogicalResult matmul_rule_impl(RewriteContext &ctx, Operation &op,
       auto new_acc_col = builder.create<tpu::MatmulOp>(
           op.getLoc(), acc_col_ty, lhs_cols[k], rhs_rolled_tile, acc_col,
           transpose_lhs, transpose_rhs, precision_attr);
-      new_acc_col->setAttr("out_layout", acc_layout_attr);
       auto new_acc_vregs = builder.create<tpu::UnrollVectorsOp>(
           op.getLoc(),
           TypeRange(ValueRange(XlaArrayToFlatArrayRef(acc_col_vregs))),