/// <X_T, Y_T, JAC_T> should be <double, double, double>
///
/// if y = Square(x), where x (and so y) are DT_COMPLEX64,
/// <X_T, Y_T, JAC_T> should be <complex64, complex64, float>
/// Note that JAC_T is always real-valued, and should be an appropriate
/// precision to host the partial derivatives for dy/dx
///
/// if y = ComplexAbs(x) where x is DT_COMPLEX64 (so y is DT_FLOAT)
/// <X_T, Y_T, JAC_T> should be <complex64, float, float>
///

    // TODO(b/188995810): DenseElementsAttr::get doesn't support complex
    // Attributes being passed, so we bail out for now. This should just be
    //   MATCH(DT_COMPLEX64, scomplex) / 2;
    //   MATCH(DT_COMPLEX128, dcomplex) / 2;
    // when DenseElementsAttr is updated.
    case DT_COMPLEX64:
    case DT_COMPLEX128:
    default:
      return -1;
  }
}

}

TEST_F(LinalgGradTest, Einsum_MatMulComplex) {
  TensorShape x_shape({2, 3});
  TensorShape y_shape({3, 3});
  Output x = Placeholder(scope_, DT_COMPLEX64, Placeholder::Shape(x_shape));
  Output y = Placeholder(scope_, DT_COMPLEX64, Placeholder::Shape(y_shape));
  auto z = Einsum(scope_, {x, y}, "ij,jk->ik");
  TensorShape z_shape({2, 3});
  TF_ASSERT_OK(scope_.status());
  float max_error;

    case DT_UINT64:
      *type = builder.getIntegerType(64, /*isSigned=*/false);
      return absl::OkStatus();
    case DT_BFLOAT16:
      *type = builder.getBF16Type();
      return absl::OkStatus();
    case DT_COMPLEX64:
      *type = mlir::ComplexType::get(builder.getF32Type());
      return absl::OkStatus();
    case DT_COMPLEX128:
      *type = mlir::ComplexType::get(builder.getF64Type());
      return absl::OkStatus();

tensorflow::DataType TflTypeToTfType(tflite::TensorType type) {
  switch (type) {
    case tflite::TensorType_BOOL:
      return tensorflow::DT_BOOL;
    case tflite::TensorType_COMPLEX64:
      return tensorflow::DT_COMPLEX64;
    case tflite::TensorType_COMPLEX128:
      return tensorflow::DT_COMPLEX128;
    case tflite::TensorType_FLOAT16:
      return tensorflow::DT_HALF;
    case tflite::TensorType_BFLOAT16:

  EXPECT_LT(max_error, 1e-10);
}

TEST(GradientCheckerTest, BasicComplex64) {
  Scope scope = Scope::NewRootScope();
  TensorShape shape({2, 4, 3});
  auto x = Placeholder(scope, DT_COMPLEX64, Placeholder::Shape(shape));
  auto y = Square(scope, x);
  float max_error;
  TF_ASSERT_OK((ComputeGradientError<complex64, complex64, float>(
      scope, {x}, {shape}, {y}, {shape}, &max_error)));

// Kernel registrations

constexpr std::array<DataType, 18> kAllXlaCpuTypes = {
    {DT_UINT8, DT_QUINT8, DT_UINT16, DT_INT8, DT_QINT8, DT_INT16, DT_INT32,
     DT_QINT32, DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE, DT_COMPLEX64,
     DT_COMPLEX128, DT_BOOL, DT_BFLOAT16, DT_INT4, DT_UINT4}};

REGISTER_XLA_LAUNCH_KERNEL(DEVICE_XLA_CPU, XlaLocalLaunchOp, kAllXlaCpuTypes);
REGISTER_XLA_COMPILE_KERNEL(DEVICE_XLA_CPU, XlaCompileOp, kAllXlaCpuTypes);

    {DT_UINT8,    DT_QUINT8,      DT_UINT16,        DT_INT8,       DT_QINT8,
     DT_INT16,    DT_INT32,       DT_QINT32,        DT_INT64,      DT_HALF,
     DT_FLOAT,    DT_DOUBLE,      DT_COMPLEX64,     DT_COMPLEX128, DT_BOOL,
     DT_BFLOAT16, DT_FLOAT8_E5M2, DT_FLOAT8_E4M3FN, DT_INT4,       DT_UINT4}};

REGISTER_XLA_LAUNCH_KERNEL(DEVICE_XLA_GPU, XlaLocalLaunchOp, kAllXlaGpuTypes);

          dim {
            size: -1
          }
        }
      }
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_HALF
        type: DT_COMPLEX64
        type: DT_COMPLEX128
      }
    }
  }
}
node {
  name: "InfeedDequeueTuple"
  op: "InfeedDequeueTuple"
  attr {
    key: "shapes"
    value {
      list {

                 mlir::dyn_cast<mlir::ComplexType>(element_type)) {
    auto etype = complex_type.getElementType();
    if (etype.isF32()) {
      tensorflow::TensorProto repr;
      repr.set_dtype(tensorflow::DT_COMPLEX64);

      tensorflow::TensorShapeProto* shape = repr.mutable_tensor_shape();
      shape->set_unknown_rank(false);
      shape->add_dim()->set_size(int64_t{1});
      std::string content;