node {
  name: "x"
  op: "VarHandleOp"
  device: "/CPU:0"
  attr {
    key: "container"
    value {
      s: "a"
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }
  attr {
    key: "shared_name"
    value {
      s: "x"
    }
  }
}
node {

    case DT_INT16:
      *type = builder.getIntegerType(16);
      return absl::OkStatus();
    case DT_INT32:
      *type = builder.getIntegerType(32);
      return absl::OkStatus();
    case DT_INT64:
      *type = builder.getIntegerType(64);
      return absl::OkStatus();
    case DT_UINT8:
      *type = builder.getIntegerType(8, /*isSigned=*/false);
      return absl::OkStatus();
    case DT_UINT16:

// CHECK-NEXT:   value {
// CHECK-NEXT:     i: 0
// CHECK-NEXT:   }
// CHECK-NEXT: }
// CHECK-NEXT: attr {
// CHECK-NEXT:   key: "Tdense"
// CHECK-NEXT:   value {
// CHECK-NEXT:     list {
// CHECK-NEXT:       type: DT_INT64
// CHECK-NEXT:     }
// CHECK-NEXT:   }
// CHECK-NEXT: }
// CHECK-NEXT: attr {
// CHECK:        key: "dense_shapes"
// CHECK-NEXT:   value {
// CHECK-NEXT:     list {
// CHECK-NEXT:       shape {

    case tflite::TensorType_INT32:
      return tensorflow::DT_INT32;
    case tflite::TensorType_UINT32:
      return tensorflow::DT_UINT32;
    case tflite::TensorType_INT64:
      return tensorflow::DT_INT64;
    case tflite::TensorType_STRING:
      return tensorflow::DT_STRING;
    case tflite::TensorType_UINT8:
      return tensorflow::DT_UINT8;
    case tflite::TensorType_UINT16:
      return tensorflow::DT_UINT16;

  attr {
    key: "ragged_value_types"
    value {
      list {
      }
    }
  }
  attr {
    key: "sparse_types"
    value {
      list {
        type: DT_STRING
        type: DT_INT64
      }
    }
  }
}
versions {
  producer: 175

    CONVERT_FLAT(DT_FLOAT, float)
    CONVERT_FLAT(DT_DOUBLE, double)
    CONVERT_FLAT(DT_INT8, int8)
    CONVERT_FLAT(DT_INT16, int16)
    CONVERT_FLAT(DT_INT32, int32)
    CONVERT_FLAT(DT_INT64, int64_t)
    CONVERT_FLAT(DT_UINT8, uint8)
    CONVERT_FLAT(DT_UINT16, uint16)
    CONVERT_FLAT(DT_UINT32, uint32)
    CONVERT_FLAT(DT_UINT64, uint64)
    CONVERT_FLAT(DT_COMPLEX64, std::complex<float>)

      {"y: T"},
      // Attr def
      {"T: {float, double, int32, int64}"},
      // Nodes
      {
          {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
          {{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
          {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
      });
}

class MlirGraphOptimizationPassTest : public Test {
 public:

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

// Kernel registrations

constexpr std::array<DataType, 20> kAllXlaGpuTypes = {
    {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}};

      OP_REQUIRES(context, context->input_type(i * 3 + 1) == DT_FLOAT,
                  absl::AbortedError("The input `max` must have float type."));
      OP_REQUIRES(
          context, context->input_type(i * 3 + 2) == DT_INT64,
          absl::AbortedError("The input `histogram` must have int64 type."));
    }
  }

  ~CalibrationStatisticsSaverOp() override {
    // Save to file during destruction so we only save it once.