// reshapes in many cases. More precisely, this function creates 1s matrix
// with appropriate shape to match with the shape of XlaDotV2 result.
// We didn't apply XlaEinsum or XlaDotV2 for this work, since it would loose
// the chance for constant folding later. We could try to add some
// postprocessing passes later to further optimize the graph after constant
// folding.

//===----------------------------------------------------------------------===//
// tf.XlaEinsum legalization
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @xlaeinsum
func.func @xlaeinsum(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> {
  // CHECK-NEXT:  mhlo.einsum

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  TF_DerivedOperandTypeAttr Tindices = TF_DerivedOperandTypeAttr<2>;
}

def TF_XlaEinsumOp : TF_Op<"XlaEinsum", [Pure, TF_NoConstantFold]> {
  let summary = [{
An op which supports basic einsum op with 2 inputs and 1 output.
  }];

  let description = [{