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.

  input_signature.tuple().emplace_back(input_tensor_spec);
  input_signature.tuple().emplace_back(input_tensor_spec);
  // Incorrect type.
  TaggedValue output_tensor_spec(unknown_shape, tensorflow::DT_INT64);
  Status s = tf_function.RegisterTrace(std::move(trace), input_signature,
                                       output_tensor_spec);
  ASSERT_TRUE(s.ok()) << s.message();
  TaggedValue args = TaggedValue::Tuple();

  ASSERT_NO_FATAL_FAILURE(VerifyConversion<int32>(
      {1, -1}, DT_INT32, mlir::IntegerType::get(&context, 32)));
  ASSERT_NO_FATAL_FAILURE(VerifyConversion<int64_t>(
      {1, -1}, DT_INT64, mlir::IntegerType::get(&context, 64)));

  ASSERT_NO_FATAL_FAILURE(VerifyConversion<uint4>(
      {static_cast<uint4>(1), static_cast<uint4>(2)}, DT_UINT4,
      mlir::IntegerType::get(

}

TEST(XlaCompilationTest, DontAutoClusterOpsProducingVariant) {
  Scope root = Scope::NewRootScope().ExitOnError();
  Output a = ops::Placeholder(root.WithOpName("test/a"), DT_INT64);
  Output b = ops::Placeholder(root.WithOpName("test/b"), DT_INT64);

  Output cast_a = ops::Cast(root.WithOpName("test/cast_a"), a, DT_INT32);
  Output cast_b = ops::Cast(root.WithOpName("test/cast_b"), b, DT_INT32);

      return DT_UINT16;
    case toco::IODataType::INT32:
      return DT_INT32;
    case toco::IODataType::UINT32:
      return DT_UINT32;
    case toco::IODataType::INT64:
      return DT_INT64;
    case toco::IODataType::UINT8:
      return DT_UINT8;
    case toco::IODataType::UINT64:
      return DT_UINT64;
    case toco::IODataType::STRING:
      return DT_STRING;
    case toco::IODataType::BOOL:

    case DT_INT16:
      return CompareTensor<int16>(tensor, expected_tensor, listener);
    case DT_INT32:
      return CompareTensor<int32>(tensor, expected_tensor, listener);
    case DT_INT64:
      return CompareTensor<int64_t>(tensor, expected_tensor, listener);
    case DT_UINT8:
      return CompareTensor<uint8>(tensor, expected_tensor, listener);
    case DT_UINT16: