}

std::string ExpectedString(const char* type) {
  const auto format_str = R"str(kernel {
  op: "TypeOp%s"
  device_type: "FakeDeviceName1"
  constraint {
    name: "T"
    allowed_values {
      list {
        type: %s
      }
    }
  }
}
)str";
  return absl::StrFormat(format_str, type, type);
}

#define TEST_KERNEL_TYPE_CONSTRAINT(tf_type, dtype)                          \

  }

  tensorflow::DataType dtype;
  auto status = tensorflow::ConvertToDataType(operand.getType(), &dtype);
  if (!status.ok()) {
    op->emitRemark() << "skipping legalization due to " << status.ToString();
    return tensorflow::XlaExpression::Invalid();
  }
  return tensorflow::XlaExpression::XlaOp(xla_op, dtype);
}

}  // namespace mhlo

        // Identity node that is forwarding the value of a frozen
        // VarhandleOp. We ensure that the dtype matches of the variable dtype.
        NodeDef* new_node = frozen_graph_def->add_node();
        *new_node = node;
        (*new_node->mutable_attr())["T"] =
            name_to_node_map.at(*handle_name)->attr().at("dtype");
        continue;
      }
    }
    // If the node isn't a variable, just copy the node as-is.

    }
    if (TFE_TensorHandleDataType(component.get()) != dtype) {
      TF_SetStatus(status, TF_INTERNAL,
                   "Components of a ParallelTensor must all have "
                   "the same dtype");
      return nullptr;
    }
  }
  return std::unique_ptr<ParallelTensor>(
      new ParallelTensor(parallel_device, std::move(components), shape, dtype));
}

    AllocatorAttributes attr;
    if (cc_input.dtype() == ::tensorflow::DT_VARIANT) {
      attr.set_on_host(true);
    }
    TF_RETURN_IF_ERROR(cc_ctx->allocate_temp(cc_input.dtype(), cc_input.shape(),
                                             cc_out, attr));

    switch (cc_input.dtype()) {
      case ::tensorflow::DT_INVALID: {
        *cc_out = Tensor(::tensorflow::DT_INVALID);

                .values();
        for (const StructuredValue& arg : args) {
          PartialTensorShape shape = arg.tensor_spec_value().shape();
          DataType dtype = arg.tensor_spec_value().dtype();
          TaggedValue tensor_spec(shape, dtype);
          input_signature.tuple().emplace_back(tensor_spec);
        }

        // Setup output signature.
        TensorSpecProto output_tensor_spec_proto =

                                bool unknown_rank, DataType dtype) {
  TensorShapeProto shape;
  shape.set_unknown_rank(unknown_rank);
  for (int64_t dim : dims) {
    shape.add_dim()->set_size(dim);
  }
  mlir::MLIRContext context;
  mlir::Builder b(&context);
  auto status_or = ConvertToMlirTensorType(shape, dtype, &b);
  std::string buf;
  llvm::raw_string_ostream os(buf);
  status_or.value().print(os);

#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include "tensorflow/core/framework/types.pb.h"

// Returns whether the given dtype is a quantization type in TensorFlow.
static bool IsQuantizationType(tensorflow::DataType dtype) {
  switch (dtype) {
    case tensorflow::DT_QINT8:
    case tensorflow::DT_QUINT8:
    case tensorflow::DT_QINT16:
    case tensorflow::DT_QUINT16:
    case tensorflow::DT_QINT32:

    }

    se::Stream* const src_compute_stream = src_xla_context->stream();
    TF_RET_CHECK(src_compute_stream != nullptr);
    TF_RET_CHECK(input->dtype() == output->dtype())
        << "input type: " << DataTypeString(input->dtype()) << " output type "
        << DataTypeString(output->dtype());
    TF_RET_CHECK(input->shape() == output->shape());
    TF_RET_CHECK(DMAHelper::CanUseDMA(input));
    auto* const src_compute_stream_impl =

func.func @main() -> (tensor<1x2xf16>, tensor<2xf16>) {
  %graph:2 = tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<1.0> : tensor<1x2xf16>} : () -> tensor<1x2xf16> loc("const1")
    %1:2 = tf_executor.island wraps "tf.Const"() {device = "", dtype = "tfdtype$DT_HALF", value = dense<[1.0, 2.0]> : tensor<2xf16>} : () -> tensor<2xf16> loc("const2")