REGISTER_KERNEL_BUILDER(Name("XlaSvd").Device(DEVICE),                       \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaDot").Device(DEVICE),                       \
                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(Name("XlaDotV2").Device(DEVICE),                     \

// CHECK: %[[XLADOT:.*]] = "tf.XlaDotV2"(%arg0, %[[SHARDED_W]]) <{dimension_numbers = "\12\01\00\0A\01\03", precision_config = ""}> {device = ""} : (tensor<?x?x?x?xf32>, tensor<512x512xf32>) -> tensor<?x?x?x?xf32>
// CHECK: %[[ORIGINAL_CAST:.*]] = "tf.Cast"(%[[XLADOT]]) <{Truncate = false}> : (tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xbf16>

                                         "XlaConv",
                                         "XlaConvV2",
                                         "XlaDequantize",
                                         "XlaDot",
                                         "XlaDotV2",
                                         "XlaDynamicSlice",
                                         "XlaDynamicUpdateSlice",

      d2.add_rhs_batch_dimensions(v);
    }
    for (auto v : d1.rhs_contracting_dimensions()) {
      d2.add_rhs_contracting_dimensions(v);
    }
  }
}

// Figure out the shape of other xladot argument for reducing contracting
// dimension.
// It must have the contracting dimensions on its shape, to reduce the
// contracting dims from the original target. In addition, to match with

      "XlaBroadcastHelper",
      "XlaCallModule",
      "XlaConcatND",
      "XlaConv",
      "XlaConvV2",
      "XlaCustomCall",
      "XlaCustomCallV2",
      "XlaDequantize",
      "XlaDot",
      "XlaDotV2",
      "XlaDynamicSlice",
      "XlaDynamicUpdateSlice",
      "XlaEinsum",
      "XlaGather",
      "XlaIf",
      "XlaKeyValueSort",
      "XlaOptimizationBarrier",

  func.return %0 : tensor<4x4x14x14x16xf32>
}

//===----------------------------------------------------------------------===//
// tf.XlaDot legalization
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: @xladot_matmul(

  if (!dims.ParseFromString(attr.getValue().str())) return {};
  return ::xla::ConvertGatherDimensionNumbers(dims, builder);
}

//===----------------------------------------------------------------------===//
// XlaDot op utilities.
//===----------------------------------------------------------------------===//

bool HasValidDotDims(StringAttr attr) {
  ::xla::DotDimensionNumbers dims;

  );

  TF_DerivedOperandTypeListAttr operand_dtypes = TF_DerivedOperandTypeListAttr<0>;
  TF_DerivedResultTypeListAttr result_dtypes = TF_DerivedResultTypeListAttr<0>;
}

def TF_XlaDotOp : TF_Op<"XlaDot", [Pure, TF_NoConstantFold]> {
  let summary = "Wraps the XLA DotGeneral operator, documented at";

  let description = [{
https://www.tensorflow.org/performance/xla/operation_semantics#dotgeneral
.