}

TEST(OStreamTest, TestTensorSpec) {
  std::stringstream stream;
  TensorSpec tensor_spec;
  tensor_spec.shape = tensorflow::PartialTensorShape({2});
  tensor_spec.dtype = tensorflow::DT_FLOAT;
  stream << tensor_spec;
  ASSERT_EQ(stream.str(), "TensorSpec(shape = [2], dtype = 1)");
}

}  // namespace impl
}  // namespace libtf

          DataType dtype = n->output_type(0);
          AddNodeAttr("dtype", dtype, &const_def);
          TensorProto value;
          value.set_dtype(dtype);
          value.mutable_tensor_shape()->add_dim()->set_size(
              shape_proto.dim_size());
          for (const auto& dim : shape_proto.dim()) {
            if (dtype == DT_INT32) {
              value.add_int_val(dim.size());

absl::StatusOr<std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
ProcessTensorSpec(mlir::TensorType type) {
  tensorflow::DataType dtype;
  TF_RETURN_IF_ERROR(
      ConvertScalarTypeToDataType(type.getElementType(), &dtype));

  if (!type.hasRank())
    return std::make_pair(dtype, tensorflow::PartialTensorShape());

  auto shape = type.getShape();
  llvm::SmallVector<int64_t, 4> dims;
  dims.assign(shape.begin(), shape.end());

node {
  name: "Placeholder"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 2
        }
        dim {
          size: 5
        }
        dim {

    Tensor* cpu_tensor, StatusCallback done) {
  DataType dtype = EncodePrimitiveTypeAsDataType(shape_.element_type()).value();
  TensorShape tensor_shape;
  Status status = XLAShapeToTensorShape(shape_, &tensor_shape);
  if (!status.ok()) {
    done(status);
    return;
  }

  *cpu_tensor = Tensor(dtype, tensor_shape);

  status = stream_->Memcpy(cpu_tensor->data(), device_memory_base_,

  // Copy XLA results to the OpOutputList.
  int output_num = 0;
  for (int i = 0, end = ctx->num_outputs(); i < end; ++i) {
    const DataType& type = compilation_result.outputs[i].type;
    VLOG(2) << "Populating output for retval " << i << " type "
            << DataTypeString(type);
    if (type == DT_VARIANT) {
      return absl::UnimplementedError(
          "Support for TensorList crossing the XLA/TF boundary "

  class AddOneModel(tf.keras.Model):

    def call(self, x):
      return x + 1

  model = AddOneModel()

  x_train = tf.constant([1, 2, 3, 4, 5], dtype=tf.float32)
  y_train = tf.constant([2, 3, 4, 5, 6], dtype=tf.float32)

  model.compile(optimizer='sgd', loss='mse')
  model.fit(x_train, y_train, epochs=1)

  path = tempdir + '/add_one_model'
  model.save(path)
  return path

func (p *Package) recordTypedefs(dtype dwarf.Type, pos token.Pos) {
	p.recordTypedefs1(dtype, pos, map[dwarf.Type]bool{})
}

func (p *Package) recordTypedefs1(dtype dwarf.Type, pos token.Pos, visited map[dwarf.Type]bool) {
	if dtype == nil {
		return
	}
	if visited[dtype] {
		return
	}
	visited[dtype] = true
	switch dt := dtype.(type) {
	case *dwarf.TypedefType:

class ConvertTensorTest : public ::testing::Test {
 protected:
  template <typename T>
  void VerifyConversion(std::initializer_list<T> values, DataType dtype,
                        mlir::Type expected_ty) {
    mlir::Builder b(expected_ty.getContext());
    Tensor tensor(dtype, TensorShape({static_cast<int64_t>(values.size())}));
    tensor.flat<T>().setValues(values);

    auto value_or = ConvertTensor(tensor, &b);

      shape_determination_fns.layout_preference_fn(
          cpu_tensor->shape(), cpu_tensor->dtype(), std::nullopt);
  TF_ASSIGN_OR_RETURN(xla::Shape shape,
                      shape_determination_fns.shape_representation_fn(
                          cpu_tensor->shape(), cpu_tensor->dtype(),
                          /*fast_mem=*/false, layout_preference));
  const xla::Layout* device_layout = &(shape.layout());