%cst_0 = "tf.Const"() {value = dense<(1.000000e+00,1.000000e+00)> : tensor<8x64x128xcomplex<f32>>} : () -> tensor<8x64x128xcomplex<f32>>
  // expected-error @+1 {{Could not legalize op: tf.XlaEinsum}}
  %1 = "tf.XlaEinsum"(%cst_0, %0) {equation = "abc,cde->abde"} : (tensor<8x64x128xcomplex<f32>>, tensor<128x32x4xcomplex<f32>>) -> tensor<8x64x32x4xcomplex<f32>>
  return %cst_0 : tensor<8x64x128xcomplex<f32>>
}

// -----

  REGISTER_KERNEL_BUILDER(Name("XlaDequantize").Device(DEVICE),                \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaEinsum").Device(DEVICE),                    \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaRngBitGenerator")                           \

                                         "XlaDotV2",
                                         "XlaDynamicSlice",
                                         "XlaDynamicUpdateSlice",
                                         "XlaEinsum",
                                         "XlaGather",
                                         "XlaIf",
                                         "XlaKeyValueSort",
                                         "XlaPad",

// 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

      "XlaConvV2",
      "XlaCustomCall",
      "XlaCustomCallV2",
      "XlaDequantize",
      "XlaDot",
      "XlaDotV2",
      "XlaDynamicSlice",
      "XlaDynamicUpdateSlice",
      "XlaEinsum",
      "XlaGather",
      "XlaIf",
      "XlaKeyValueSort",
      "XlaOptimizationBarrier",
      "XlaPad",
      "XlaRecv",
      "XlaReduce",
      "XlaReducePrecision",

  );

  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 = [{