Function TEST

tensorflow/core/util/tensor_slice_writer_test.cc:73–123

A simple test on writing a few tensor slices TODO(yangke): refactor into smaller tests: will do as we add more stuff to the writer.

// TODO(yangke): refactor into smaller tests: will do as we add more stuff to
// the writer.
TEST(TensorSliceWriteTest, SimpleWrite) {
  const string filename = io::JoinPath(testing::TmpDir(), "checkpoint");

  TensorSliceWriter writer(filename, CreateTableTensorSliceBuilder);

  // Add some int32 tensor slices
  {
    TensorShape shape({5, 10});
    TensorSlice slice = TensorSlice::ParseOrDie("-:0,1");
    const int32 data[] = {0, 1, 2, 3, 4};
    TF_CHECK_OK(writer.Add("test", shape, slice, data));
  }

  // Two slices share the same tensor name
  {
    TensorShape shape({5, 10});
    TensorSlice slice = TensorSlice::ParseOrDie("-:3,1");
    const int32 data[] = {10, 11, 12, 13, 14};
    TF_CHECK_OK(writer.Add("test", shape, slice, data));
  }

  // Another slice from a different float tensor -- it has a different name and
  // should be inserted in front of the previous tensor
  {
    TensorShape shape({3, 2});
    TensorSlice slice = TensorSlice::ParseOrDie("-:-");
    const float data[] = {1.2, 1.3, 1.4, 2.1, 2.2, 2.3};
    TF_CHECK_OK(writer.Add("AA", shape, slice, data));
  }

  // A slice with int64 data
  {
    TensorShape shape({5, 10});
    TensorSlice slice = TensorSlice::ParseOrDie("-:3,1");
    const int64 data[] = {10, 11, 12, 13, 14};
    TF_CHECK_OK(writer.Add("int64", shape, slice, data));
  }

  // A slice with int16 data
  {
    TensorShape shape({5, 10});
    TensorSlice slice = TensorSlice::ParseOrDie("-:3,1");
    const int16 data[] = {10, 11, 12, 13, 14};
    TF_CHECK_OK(writer.Add("int16", shape, slice, data));
  }

  TF_CHECK_OK(writer.Finish());

  // Now we examine the checkpoint file manually.
  TensorSliceWriteTestHelper::CheckEntries(filename);
}

}  // namespace