mlir::IntegerType::get(
          &context, 64, mlir::IntegerType::SignednessSemantics::Unsigned)));

  ASSERT_NO_FATAL_FAILURE(VerifyConversion<std::complex<float>>(
      {{0.0, 1.0}, {1.0, 0.0}}, DT_COMPLEX64,
      mlir::ComplexType::get(mlir::FloatType::getF32(&context))));
  ASSERT_NO_FATAL_FAILURE(VerifyConversion<std::complex<double>>(
      {{0.0, 1.0}, {1.0, 0.0}}, DT_COMPLEX128,

      return DT_UINT64;
    case toco::IODataType::STRING:
      return DT_STRING;
    case toco::IODataType::BOOL:
      return DT_BOOL;
    case toco::IODataType::COMPLEX64:
      return DT_COMPLEX64;
    case toco::IODataType::COMPLEX128:
      return DT_COMPLEX128;
    case toco::IODataType::RESOURCE:
      return DT_RESOURCE;
    case toco::IODataType::VARIANT:
      return DT_VARIANT;
    default:

              type: DT_DOUBLE
              type: DT_UINT8
              type: DT_INT8
              type: DT_INT16
              type: DT_INT32
              type: DT_INT64
              type: DT_COMPLEX64
              type: DT_COMPLEX128
              type: DT_STRING
            }
          }
        }
      }
      op {
        name: "Assign"
        input_arg {
          name: "ref"

// Conjugate helper function returns the conjugate of an Output if it
// is complex valued.
Output ConjugateHelper(const Scope& scope, const Output& out) {
  DataType dtype = out.type();
  if (dtype == DT_COMPLEX64 || dtype == DT_COMPLEX128) {
    return Conj(scope, out);
  } else {
    return out;
  }
}

// TODO(andydavis) Add control dependencies to gradient functions (as needed).

TEST_KERNEL_TYPE_CONSTRAINT(TF_UINT16, DT_UINT16);
TEST_KERNEL_TYPE_CONSTRAINT(TF_UINT8, DT_UINT8);
TEST_KERNEL_TYPE_CONSTRAINT(TF_INT8, DT_INT8);
TEST_KERNEL_TYPE_CONSTRAINT(TF_INT32, DT_INT32);
TEST_KERNEL_TYPE_CONSTRAINT(TF_COMPLEX64, DT_COMPLEX64);
TEST_KERNEL_TYPE_CONSTRAINT(TF_COMPLEX128, DT_COMPLEX128);
TEST_KERNEL_TYPE_CONSTRAINT(TF_QINT8, DT_QINT8);
TEST_KERNEL_TYPE_CONSTRAINT(TF_QUINT8, DT_QUINT8);
TEST_KERNEL_TYPE_CONSTRAINT(TF_QINT32, DT_QINT32);

};

// Template instantiations here

inline bool IsDtypeTrainable(DataType dtype) {
  switch (dtype) {
    case DT_HALF:
    case DT_BFLOAT16:
    case DT_FLOAT:
    case DT_DOUBLE:
    case DT_COMPLEX64:
    case DT_COMPLEX128:
    case DT_RESOURCE:
    case DT_VARIANT:
      return true;
    case DT_QINT8:
    case DT_QINT16:
    case DT_QINT32:
    case DT_QUINT8:
    case DT_QUINT16:

  EXPECT_EQ(TF_INT8, static_cast<TF_DataType>(tensorflow::DT_INT8));
  EXPECT_EQ(TF_STRING, static_cast<TF_DataType>(tensorflow::DT_STRING));
  EXPECT_EQ(TF_COMPLEX64, static_cast<TF_DataType>(tensorflow::DT_COMPLEX64));
  EXPECT_EQ(TF_COMPLEX, TF_COMPLEX64);
  EXPECT_EQ(TF_INT64, static_cast<TF_DataType>(tensorflow::DT_INT64));
  EXPECT_EQ(TF_BOOL, static_cast<TF_DataType>(tensorflow::DT_BOOL));

          dtype == tensorflow::DT_HALF || dtype == tensorflow::DT_BFLOAT16 ||
          dtype == tensorflow::DT_FLOAT || dtype == tensorflow::DT_DOUBLE ||
          dtype == tensorflow::DT_COMPLEX64 ||
          dtype == tensorflow::DT_COMPLEX128 || dtype == tensorflow::DT_BOOL) {
        return {};
      }
    }
  }

  // Permutations defined by constant operations.