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:

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

                 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;