Jay Conrod <******@****.***> 1581363722 -0500

}

//===----------------------------------------------------------------------===//
// Einsum.
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @einsum
func.func @einsum(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> {
  // CHECK:  mhlo.einsum
  %0 = "tf.Einsum"(%arg0, %arg1) {equation = "ab,bc->ac"} : (tensor<2x3xf32>, tensor<3x4xf32>) -> tensor<2x4xf32>

  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
  // pattern's benefit can be changed back to one as well as the fallback
  // lowering pattern for the op can be removed.