};
  broadcast_one_side(lhs, lhs_type, lhs_splitted.getTail(), out_lhs);
  broadcast_one_side(rhs, rhs_type, rhs_splitted.getTail(), out_rhs);
}

class ConvertBatchMatMulV2Op : public OpRewritePattern<TF::BatchMatMulV2Op> {
 public:
  // TODO(hinsu): Legalize this op to Einsum op. HLO Einsum op needs to be moved
  // to CHLO and it is missing legalization to MHLO. Once that is done, this