Commit message: Add BuildTensorSlice for building from unvalidated TensorSliceProtos.

This avoids several sources of crashes and undefined behavior when loading
invalid checkpoints.

PiperOrigin-RevId: 392785704
Change-Id: Icd9713c768b882f3b58b427eddac376060696833
CVE: CVE-2021-41203
CWE: CWE-345
Vulnerability description: TensorFlow is an open source platform for machine learning. In affected versions an attacker can trigger undefined behavior, integer overflows, segfaults and `CHECK`-fail crashes if they can change saved checkpoints from outside of TensorFlow. This is because the checkpoints loading infrastructure is missing validation for invalid file formats. The fixes will be included in TensorFlow 2.7.0. We will also cherrypick these commits on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.

Project: tensorflow
File: tensor_slice_reader.cc

TensorSliceReader::TensorSliceReader(const string& filepattern,
                                     OpenTableFunction open_function)
    : TensorSliceReader(filepattern, std::move(open_function), kLoadAllShards) {
}

Project: tensorflow
File: tensor_slice_reader.cc

TensorSliceReader::GetVariableToDataTypeMap() const {
  VarToDataTypeMap name_to_dtype;
  if (status().ok()) {
    for (auto& e : Tensors()) {
      name_to_dtype[e.first] = e.second->type();
    }
  }
  return name_to_dtype;
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(CheckpointVersionTest, MinConsumer) {
  VersionDef versions;
  versions.set_producer(TF_CHECKPOINT_VERSION + 1);
  versions.set_min_consumer(TF_CHECKPOINT_VERSION + 1);
  VersionTest(
      versions,
      strings::StrCat("Checkpoint min consumer version ",
                      TF_CHECKPOINT_VERSION + 1, " above current version ",
                      TF_CHECKPOINT_VERSION, " for TensorFlow"));
}

Project: tensorflow
File: tensor_slice_reader.cc

TensorSliceReader::Table::~Table() {}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, Deserialization) {
  // Serialization of
  //     extent { length: 5 }
  //     extent { start: 0 length: 10 }
  //     extent { start: 14 length: 1 }
  //     extent { start: 1 }
  //     extent { }
  // in proto2 and proto3:
  const char pb2[] =
      "\x0A\x02\x10\x05\x0A\x04\x08\x00"
      "\x10\x0A\x0A\x04\x08\x0E\x10\x01\x0A\x02\x08\x01\x0A\x00";
  const char pb3[] =
      "\x0A\x02\x10\x05\x0A\x02"
      "\x10\x0A\x0A\x04\x08\x0E\x10\x01\x0A\x02\x08\x01\x0A\x00";
  // (The difference is that in the proto3 version, "start: 0" isn't included
  // since 0 is start's default value.)

  TensorSliceProto proto2;
  ASSERT_TRUE(proto2.ParseFromArray(pb2, sizeof(pb2) - 1));
  TensorSlice ts2(proto2);

  TensorSliceProto proto3;
  ASSERT_TRUE(proto3.ParseFromArray(pb3, sizeof(pb3) - 1));
  TensorSlice ts3(proto3);

  // Both serializations should be interpreted the same.
  EXPECT_EQ("0,5:0,10:14,1:1,-1:-", ts2.DebugString());
  EXPECT_EQ("0,5:0,10:14,1:1,-1:-", ts3.DebugString());
}

Project: tensorflow
File: tensor_slice_reader_test.cc

void SimpleFloatHelper(
    const TensorSliceWriter::CreateBuilderFunction& create_function,
    TensorSliceReader::OpenTableFunction open_function) {
  const string fname_base = io::JoinPath(testing::TmpDir(), "float_checkpoint");

  TensorShape shape({4, 5});

  // File #0 contains a slice that is the top two rows:
  //
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    const string fname = strings::StrCat(fname_base, "_0");
    TensorSliceWriter writer(fname, create_function);
    const float data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
    TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
    TF_CHECK_OK(writer.Add("test", shape, slice, data));
    TF_CHECK_OK(writer.Finish());
  }

  // File #1 contains two slices:
  //
  // slice #0 is the bottom left corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //  10  11  12   .   .
  //  15  16  17   .   .
  //
  // slice #1 is the bottom right corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .  18  19
  {
    const string fname = strings::StrCat(fname_base, "_1");
    TensorSliceWriter writer(fname, create_function);
    // slice #0
    {
      const float data[] = {10, 11, 12, 15, 16, 17};
      TensorSlice slice = TensorSlice::ParseOrDie("2,2:0,3");
      TF_CHECK_OK(writer.Add("test", shape, slice, data));
    }
    // slice #1
    {
      const float data[] = {18, 19};
      TensorSlice slice = TensorSlice::ParseOrDie("3,1:3,2");
      TF_CHECK_OK(writer.Add("test", shape, slice, data));
    }
    TF_CHECK_OK(writer.Finish());
  }

  // Notice that we leave a hole in the tensor
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   . (13) (14)
  //   .   .   .   .   .

  // Now we need to read the tensor slices
  const string filepattern = strings::StrCat(fname_base, "_*");
  TensorSliceReader reader(filepattern, std::move(open_function));
  TF_EXPECT_OK(reader.status());
  EXPECT_EQ(2, reader.num_files());

  // We query some of the tensors
  {
    TensorShape shape;
    DataType type;
    EXPECT_TRUE(reader.HasTensor("test", &shape, &type));
    EXPECT_EQ("[4,5]", shape.DebugString());
    EXPECT_EQ(DT_FLOAT, type);
    EXPECT_FALSE(reader.HasTensor("don't exist", nullptr, nullptr));
  }

  // Now we query some slices
  //
  // Slice #1 is an exact match
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("0,2:-");
    float expected[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
    float results[10];
    EXPECT_TRUE(reader.CopySliceData("test", s, results));
    for (int i = 0; i < 10; ++i) {
      EXPECT_EQ(expected[i], results[i]);
    }
  }

  // Slice #2 is a subset match
  //   .   .   .   .   .
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,1:-");
    float expected[] = {5, 6, 7, 8, 9};
    float results[5];
    EXPECT_TRUE(reader.CopySliceData("test", s, results));
    for (int i = 0; i < 5; ++i) {
      EXPECT_EQ(expected[i], results[i]);
    }
  }

  // Slice #4 includes the hole and so there is no match
  //   .   .   .   .   .
  //   .   .   7   8   9
  //   .   .  12  13  14
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,2:2,3");
    float results[6];
    EXPECT_FALSE(reader.CopySliceData("test", s, results));
  }
}

Project: tensorflow
File: tensor_slice_reader.cc

const string TensorSliceReader::DebugString() const {
  string shape_str;
  if (status().ok()) {
    for (const auto& e : Tensors()) {
      strings::StrAppend(&shape_str, e.first, " (",
                         DataType_Name(e.second->type()), ") ",
                         e.second->shape().DebugString());
      // Indicates if a tensor has more than 1 slice (i.e., it's partitioned).
      const int num_slices = e.second->Slices().size();
      if (num_slices > 1) {
        strings::StrAppend(&shape_str, ", ", num_slices, " slices");
      }
      strings::StrAppend(&shape_str, "\n");
    }
  }
  return shape_str;
}

Project: tensorflow
File: tensor_slice_reader.cc

  bool Get(const string& key, string* value) override {
    std::unique_ptr<table::Iterator> iter(table_->NewIterator());
    iter->Seek(key);
    if (iter->Valid() && iter->key() == key) {
      StringPiece v = iter->value();
      value->assign(v.data(), v.size());
      return true;
    } else {
      return false;
    }
  }

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(TensorSliceReaderTest, NegativeTensorShapeDimension) {
  const string fname =
      io::JoinPath(testing::TmpDir(), "negative_dim_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TF_CHECK_OK(writer.Add("test", TensorShape({4, 5}),
                         TensorSlice::ParseOrDie("0,2:-"), data));
  TF_CHECK_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
    if (sts.has_meta()) {
      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
        for (auto& dim : *tensor.mutable_shape()->mutable_dim()) {
          dim.set_size(-dim.size());
        }
      }
    }
    return sts.SerializeAsString();
  });

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  // The negative dimension should cause loading to fail.
  EXPECT_FALSE(reader.status().ok());
}

Project: tensorflow
File: tensor_slice_reader_test.cc

static void VersionTest(const VersionDef& versions, const string& error) {
  const string path = io::JoinPath(testing::TmpDir(), "checkpoint");

  {
    // Prepare an empty checkpoint with some version information
    SavedTensorSlices sts;
    *sts.mutable_meta()->mutable_versions() = versions;
    string contents;
    EXPECT_TRUE(sts.SerializeToString(&contents));

    // Write it to disk
    TensorSliceWriter::Builder* builder;
    TF_ASSERT_OK(CreateTableTensorSliceBuilder(path, &builder));
    builder->Add(kSavedTensorSlicesKey, contents);
    int64_t file_size;
    TF_EXPECT_OK(builder->Finish(&file_size));
    delete builder;
  }

  // Read it back in and verify that we get the expected error
  TensorSliceReader reader(path, OpenTableTensorSliceReader);
  EXPECT_TRUE(reader.status().code() == error::INVALID_ARGUMENT &&
              absl::StartsWith(reader.status().error_message(), error))
      << "Expected error starting with '" << errors::InvalidArgument(error)
      << "', got '" << reader.status() << "'";
}

Project: tensorflow
File: tensor_slice.cc

Status TensorSlice::BuildTensorSlice(const TensorSliceProto& proto,
                                     TensorSlice* output) {
  output->Clear();
  output->starts_.reserve(proto.extent_size());
  output->lengths_.reserve(proto.extent_size());
  for (const auto& e : proto.extent()) {
    int64_t l = GetExtentLength(e);
    if (e.start() != 0 || l != kFullExtent) {
      if (e.start() < 0 || l <= 0) {
        return errors::InvalidArgument(
            "Expected non-negative start and positive length but got start = ",
            e.start(), ", length = ", l, ": extent = ", e.ShortDebugString());
      }
      // Calculating the extent end must not cause signed integer overflow.
      if (static_cast<uint64_t>(e.start()) + static_cast<uint64_t>(e.length()) >
          std::numeric_limits<int64_t>::max()) {
        return errors::InvalidArgument(
            "Extent end exceeds the maximum possible size: extent = ",
            e.ShortDebugString());
      }
    }
    output->starts_.push_back(e.start());
    output->lengths_.push_back(l);
  }

  return Status::OK();
}

Project: tensorflow
File: tensor_slice.cc

Status TensorSlice::SliceTensorShape(const TensorShape& shape,
                                     TensorShape* result_shape) const {
  result_shape->Clear();
  // Mismatching ranks: we can't apply the slice at all.
  if (shape.dims() != dims()) {
    return errors::Internal("Mismatching ranks: shape = ", shape.DebugString(),
                            ", slice = ", DebugString());
  }
  for (int d = 0; d < dims(); ++d) {
    if (IsFullAt(d)) {
      result_shape->AddDim(shape.dim_size(d));
    } else {
      // Check if the extent applies to the dimension
      if (end(d) <= shape.dim_size(d)) {
        // Yes: the end is within the range of the dim -- we adjust the result
        // shape so that its size along this dimension is the length of the
        // slice.
        result_shape->AddDim(length(d));
      } else {
        // The extent doesn't apply to the dimension
        result_shape->Clear();
        return errors::Internal("Extent in dimension ", d,
                                " out of bounds: shape = ", shape.DebugString(),
                                ", slice = ", DebugString());
      }
    }
  }
  // If we are here, we have successfully applied the shape.
  return Status::OK();
}

Project: tensorflow
File: tensor_slice.cc

Status TensorSlice::Parse(const string& str, TensorSlice* slice) {
  std::vector<string> items = str_util::Split(str, ':', str_util::SkipEmpty());
  slice->starts_.reserve(items.size());
  slice->lengths_.reserve(items.size());
  for (const string& x : items) {
    int64_t s, l;
    if (x == "-") {
      // "everything"
      s = 0;
      l = kFullExtent;
    } else {
      std::vector<string> sl = str_util::Split(x, ',', str_util::SkipEmpty());
      if (sl.size() != 2 || !strings::safe_strto64(sl[0], &s) ||
          !strings::safe_strto64(sl[1], &l)) {
        return errors::InvalidArgument(
            "Expected a pair of numbers or '-' "
            "but got '",
            x, "': string = ", str);
      }
      if (s < 0 || l <= 0) {
        return errors::InvalidArgument(
            "Expected non-negative start and "
            "positive length but got start = ",
            s, ", length = ", l, ": string = ", str);
      }
    }
    slice->starts_.push_back(s);
    slice->lengths_.push_back(l);
  }

  return Status::OK();
}

Project: tensorflow
File: tensor_slice.cc

bool TensorSlice::operator==(const TensorSlice& other) const {
  return dims() == other.dims() && starts_ == other.starts_ &&
         lengths_ == other.lengths_;
}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, UpdateToCover) {
  // [2:4, :, 3:]
  TensorSlice s({{2, 2}, {0, -1}, {3, 7}});
  // [:, 1:4, 2:4]
  TensorSlice other({{0, -1}, {1, 3}, {2, 2}});

  s.UpdateToCover(other);
  // [:, :, 2:]
  EXPECT_EQ("-:-:2,8", s.DebugString());
}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, BuildTensorSlice) {
  TensorSliceProto proto;
  TensorSlice({{0, -1}, {0, 10}, {14, 1}}).AsProto(&proto);
  TensorSlice s;

  // Successful building.
  {
    TF_ASSERT_OK(TensorSlice::BuildTensorSlice(proto, &s));
    EXPECT_EQ("-:0,10:14,1", s.DebugString());
  }

  // Failed building due to negative extent start.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(0)->set_start(-1);
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to negative extent length.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->set_length(-1);
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to missing extent length.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->clear_length();
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }

  // Failed building due to extent end overflowing.
  {
    TensorSliceProto invalid_proto = proto;
    invalid_proto.mutable_extent(2)->set_length(
        std::numeric_limits<int64_t>::max());
    EXPECT_FALSE(TensorSlice::BuildTensorSlice(invalid_proto, &s).ok());
  }
}

Project: tensorflow
File: tensor_slice.cc

TensorSlice::TensorSlice(const TensorSliceProto& proto) {
  starts_.reserve(proto.extent_size());
  lengths_.reserve(proto.extent_size());
  for (const auto& e : proto.extent()) {
    starts_.push_back(e.start());
    lengths_.push_back(GetExtentLength(e));
  }
}

Project: tensorflow
File: tensor_slice_reader.cc

TensorSliceReader::~TensorSliceReader() {
  for (auto& temp : tensors_) {
    delete temp.second;
  }
  tensors_.clear();
}

Project: tensorflow
File: tensor_slice_reader.cc

TensorSliceReader::VarToShapeMap TensorSliceReader::GetVariableToShapeMap()
    const {
  VarToShapeMap name_to_shape;
  if (status().ok()) {
    for (auto& e : Tensors()) {
      name_to_shape[e.first] = e.second->shape();
    }
  }
  return name_to_shape;
}

Project: tensorflow
File: tensor_slice_reader_test.cc

void MutateSavedTensorSlices(
    const std::string& fname,
    const std::function<std::string(SavedTensorSlices)>& mutator) {
  table::Options options;
  options.compression = table::kNoCompression;

  // Read all entres from the table.
  std::vector<std::pair<std::string, std::string>> entries;
  {
    std::unique_ptr<RandomAccessFile> file;
    TF_CHECK_OK(Env::Default()->NewRandomAccessFile(fname, &file));
    uint64 file_size;
    TF_CHECK_OK(Env::Default()->GetFileSize(fname, &file_size));
    table::Table* t;
    TF_CHECK_OK(table::Table::Open(options, file.get(), file_size, &t));
    std::unique_ptr<table::Table> table(t);
    std::unique_ptr<table::Iterator> it(table->NewIterator());
    for (it->Seek(""); it->Valid(); it->Next()) {
      entries.emplace_back(it->key(), it->value());
    }
    TF_CHECK_OK(it->status());
  }

  // Rewrite the table, mutating each value.
  {
    std::unique_ptr<WritableFile> file;
    TF_CHECK_OK(Env::Default()->NewWritableFile(fname, &file));
    table::TableBuilder builder(options, file.get());
    for (const auto& entry : entries) {
      SavedTensorSlices sts;
      CHECK(sts.ParseFromString(entry.second));
      builder.Add(entry.first, mutator(std::move(sts)));
    }
    TF_CHECK_OK(builder.Finish());
    TF_CHECK_OK(file->Close());
  }
}

Project: tensorflow
File: tensor_slice.cc

void TensorSlice::SetFullSlice(int dim) {
  Clear();
  starts_.reserve(dim);
  lengths_.reserve(dim);
  for (int d = 0; d < dim; ++d) {
    starts_.push_back(0);
    lengths_.push_back(kFullExtent);
  }
}

Project: tensorflow
File: tensor_slice.cc

TensorSlice::TensorSlice(int dim) { SetFullSlice(dim); }

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, Equality) {
  {  // Dims are different.
    TensorSlice slice1(3);
    TensorSlice slice2(2);
    EXPECT_TRUE(slice1 != slice2);
    EXPECT_TRUE(slice2 != slice1);
  }
  {  // Both are 3-dim full slices.
    TensorSlice slice1(3);
    TensorSlice slice2({{0, -1}, {0, -1}, {0, -1}});
    EXPECT_TRUE(slice1 == slice2);
    EXPECT_TRUE(slice2 == slice1);
  }
  {  // Differs in one dimension.
    TensorSlice slice1(3);
    TensorSlice slice2({{0, -1}, {0, 1}, {0, -1}});
    EXPECT_TRUE(slice1 != slice2);
    EXPECT_TRUE(slice2 != slice1);
  }
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(TensorSliceReaderTest, InvalidTensorSlice) {
  const string fname =
      io::JoinPath(testing::TmpDir(), "invalid_slice_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TF_CHECK_OK(writer.Add("test", TensorShape({4, 5}),
                         TensorSlice::ParseOrDie("0,2:-"), data));
  TF_CHECK_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
    if (sts.has_meta()) {
      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
        tensor.mutable_slice(0)->mutable_extent(0)->set_length(-10);
      }
    }
    return sts.SerializeAsString();
  });

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  // The negative exent length should cause loading to fail.
  EXPECT_FALSE(reader.status().ok());
}

Project: tensorflow
File: tensor_slice.cc

bool TensorSlice::IsFull() const {
  for (int d = 0; d < dims(); ++d) {
    if (!IsFullAt(d)) return false;
  }
  return true;
}

Project: tensorflow
File: tensor_slice.cc

void TensorSlice::Clear() {
  starts_.clear();
  lengths_.clear();
}

Project: tensorflow
File: tensor_slice.cc

void TensorSlice::AsProto(TensorSliceProto* proto) const {
  for (int d = 0; d < dims(); ++d) {
    TensorSliceProto::Extent* e = proto->add_extent();
    // We only need to record the explicit slice for non-full slices
    if (!IsFullAt(d)) {
      e->set_start(starts_[d]);
      e->set_length(lengths_[d]);
    }
  }
}

Project: tensorflow
File: tensor_slice_reader.cc

Status TensorSliceReader::GetTensor(
    const string& name, std::unique_ptr<tensorflow::Tensor>* out_tensor) const {
  DataType type;
  TensorShape shape;
  TensorSlice slice;
  {
    mutex_lock l(mu_);
    const TensorSliceSet* tss = gtl::FindPtrOrNull(tensors_, name);
    if (tss == nullptr) {
      return errors::NotFound(name, " not found in checkpoint file");
    }

    if (tss->Slices().size() > 1) {
      // TODO(sherrym): Support multi-slice checkpoints.
      return errors::Unimplemented("Sliced checkpoints are not supported");
    }

    type = tss->type();
    shape = tss->shape();
    slice = tss->Slices().begin()->second.slice;
  }

  std::unique_ptr<tensorflow::Tensor> t(new tensorflow::Tensor);
  Status s = tensorflow::Tensor::BuildTensor(type, shape, t.get());
  if (!s.ok()) return s;
  bool success = false;

#define READER_COPY(dt)                                                  \
  case dt:                                                               \
    success = CopySliceData(name, slice,                                 \
                            t->flat<EnumToDataType<dt>::Type>().data()); \
    break;

  switch (type) {
    READER_COPY(DT_FLOAT);
    READER_COPY(DT_DOUBLE);
    READER_COPY(DT_INT32);
    READER_COPY(DT_UINT8);
    READER_COPY(DT_INT16);
    READER_COPY(DT_INT8);
    READER_COPY(DT_INT64);
    READER_COPY(DT_STRING);
    default:
      return errors::Unimplemented("Data type not supported");
  }
#undef READER_COPY

  if (!success) {
    return errors::NotFound(name, " not found in checkpoint file");
  }
  std::swap(*out_tensor, t);

  return Status::OK();
}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, ComputeRelative) {
  // Easy case: base is "everything"
  {
    TensorSlice base = TensorSlice::ParseOrDie("-:-:-:-");
    TensorSlice sub = TensorSlice::ParseOrDie("-:1,2:-:3,4");
    TensorSlice relative;
    base.ComputeRelative(sub, &relative);
    EXPECT_EQ("-:1,2:-:3,4", relative.DebugString());
  }

  // A slightly more complicated case
  {
    TensorSlice base = TensorSlice::ParseOrDie("1,2:3,4:-:5,1");
    TensorSlice sub = TensorSlice::ParseOrDie("1,1:4,2:3,3:5,1");
    TensorSlice relative;
    base.ComputeRelative(sub, &relative);
    EXPECT_EQ("0,1:1,2:3,3:0,1", relative.DebugString());
  }
}

Project: tensorflow
File: tensor_slice_reader.cc

bool TensorSliceReader::HasTensor(const string& name, TensorShape* shape,
                                  DataType* type) const {
  mutex_lock l(mu_);
  const TensorSliceSet* tss = gtl::FindPtrOrNull(tensors_, name);
  if (!tss && !all_shards_loaded_) {
    VLOG(1) << "Did not find tensor in preferred shard, loading all shards: "
            << name;
    LoadAllShards();
    tss = gtl::FindPtrOrNull(tensors_, name);
  }
  if (tss) {
    if (shape) {
      *shape = tss->shape();
    }
    if (type) {
      *type = tss->type();
    }
    return true;
  } else {
    return false;
  }
}

Project: tensorflow
File: tensor_slice_reader.cc

TensorSliceReader::TensorSliceReader(const string& filepattern)
    : TensorSliceReader(filepattern, OpenTableTensorSliceReader,
                        kLoadAllShards) {}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, Serialization) {
  // Serialization
  {
    TensorSlice s({{0, -1}, {0, 10}, {14, 1}, {0, -1}});
    EXPECT_EQ("-:0,10:14,1:-", s.DebugString());
    EXPECT_TRUE(!s.IsFull());
  }

  {
    TensorSliceProto proto;
    // Define ptxt outside ASSERT_TRUE call to work around bug in some
    // versions of gcc that breaks when you use raw string literals
    // inside macro expansions.
    //   See: https://gcc.gnu.org/bugzilla/show_bug.cgi?id=55971
    const char* ptxt = R"pb(
      extent {}
      extent { start: 0 length: 10 }
      extent { start: 14 length: 1 }
      extent {}
    )pb";
    ASSERT_TRUE(protobuf::TextFormat::ParseFromString(ptxt, &proto));
    TensorSlice s(proto);
    EXPECT_EQ("-:0,10:14,1:-", s.DebugString());
    EXPECT_TRUE(!s.IsFull());
  }

  // Parsing
  {
    TensorSlice s = TensorSlice::ParseOrDie("-:-:1,3:4,5");
    TensorSliceProto proto;
    s.AsProto(&proto);
    EXPECT_EQ(
        "extent { } "
        "extent { } "
        "extent { start: 1 length: 3 } "
        "extent { start: 4 length: 5 }",
        proto.ShortDebugString());
    EXPECT_TRUE(!s.IsFull());
  }

  // Failed parsing
  {
    TensorSlice slice;
    Status s = TensorSlice::Parse("-:-:1,3:4:5", &slice);
    EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
    EXPECT_TRUE(
        absl::StrContains(s.error_message(),
                          "Expected a pair of numbers or '-' but got '4': "
                          "string = -:-:1,3:4:5"));
  }
  {
    TensorSlice slice;
    Status s = TensorSlice::Parse("-:-1,3", &slice);
    EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
    EXPECT_TRUE(absl::StrContains(
        s.error_message(),
        "Expected non-negative start and positive length but got "
        "start = -1, length = 3: string = -:-1,3"));
  }

  // int64 parsing
  {
    TensorSlice s =
        TensorSlice::ParseOrDie("9223372036854775807,9223372036854775807");
    TensorSliceProto proto;
    s.AsProto(&proto);
    EXPECT_EQ(
        "extent { start: 9223372036854775807 length: 9223372036854775807 }",
        proto.ShortDebugString());
    EXPECT_TRUE(!s.IsFull());
  }

  // int64 parsing failure
  {
    TensorSlice slice;
    Status s =
        TensorSlice::Parse("19223372036854775808,19223372036854775808", &slice);
    EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
    EXPECT_TRUE(absl::StrContains(
        s.error_message(),
        "Expected a pair of numbers or '-' but got "
        "'19223372036854775808,19223372036854775808': string = "
        "19223372036854775808,19223372036854775808"));
  }
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(TensorSliceReaderTest, SimpleFloat) {
  SimpleFloatHelper(CreateTableTensorSliceBuilder, OpenTableTensorSliceReader);
}

Project: tensorflow
File: tensor_slice.cc

bool TensorSlice::HasExtentLength(const TensorSliceProto::Extent& extent) {
  return extent.has_length_case() == TensorSliceProto::Extent::kLength;
}

Project: tensorflow
File: tensor_slice_reader.cc

  explicit TensorSliceReaderTable(RandomAccessFile* f, table::Table* t)
      : file_(f), table_(t) {}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, IsFull) {
  TensorSlice slice(3);
  EXPECT_TRUE(slice.IsFull());

  TensorSlice slice2({{0, -1}});
  EXPECT_TRUE(slice2.IsFull());

  TensorSlice slice3({{0, -1}, {0, -1}, {14, 1}});
  EXPECT_TRUE(!slice3.IsFull());
}

Project: tensorflow
File: tensor_slice_reader.cc

Status OpenTableTensorSliceReader(const string& fname,
                                  TensorSliceReader::Table** result) {
  *result = nullptr;
  Env* env = Env::Default();
  std::unique_ptr<RandomAccessFile> f;
  Status s = env->NewRandomAccessFile(fname, &f);
  if (s.ok()) {
    uint64 file_size;
    s = env->GetFileSize(fname, &file_size);
    if (s.ok()) {
      table::Options options;
      table::Table* table;
      s = table::Table::Open(options, f.get(), file_size, &table);
      if (s.ok()) {
        *result = new TensorSliceReaderTable(f.release(), table);
        return Status::OK();
      } else {
        s = Status(s.code(),
                   strings::StrCat(s.error_message(),
                                   ": perhaps your file is in a different "
                                   "file format and you need to use a "
                                   "different restore operator?"));
      }
    }
  }
  LOG(WARNING) << "Could not open " << fname << ": " << s;
  return s;
}

Project: tensorflow
File: tensor_slice_reader_test.cc

void SimpleIntXHelper(
    const TensorSliceWriter::CreateBuilderFunction& create_function,
    TensorSliceReader::OpenTableFunction open_function,
    const string& checkpoint_file) {
  const string fname_base = io::JoinPath(testing::TmpDir(), checkpoint_file);

  TensorShape shape({4, 5});

  // File #0 contains a slice that is the top two rows:
  //
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    const string fname = strings::StrCat(fname_base, "_0");
    TensorSliceWriter writer(fname, create_function);
    const T data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
    TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
    TF_CHECK_OK(writer.Add("test", shape, slice, data));
    TF_CHECK_OK(writer.Finish());
  }

  // File #1 contains two slices:
  //
  // slice #0 is the bottom left corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //  10  11  12   .   .
  //  15  16  17   .   .
  //
  // slice #1 is the bottom right corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .  18  19
  {
    const string fname = strings::StrCat(fname_base, "_1");
    TensorSliceWriter writer(fname, create_function);
    // slice #0
    {
      const T data[] = {10, 11, 12, 15, 16, 17};
      TensorSlice slice = TensorSlice::ParseOrDie("2,2:0,3");
      TF_CHECK_OK(writer.Add("test", shape, slice, data));
    }
    // slice #1
    {
      const T data[] = {18, 19};
      TensorSlice slice = TensorSlice::ParseOrDie("3,1:3,2");
      TF_CHECK_OK(writer.Add("test", shape, slice, data));
    }
    TF_CHECK_OK(writer.Finish());
  }

  // Notice that we leave a hole in the tensor
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   . (13) (14)
  //   .   .   .   .   .

  // Now we need to read the tensor slices
  const string filepattern = strings::StrCat(fname_base, "_*");
  TensorSliceReader reader(filepattern, std::move(open_function));
  TF_EXPECT_OK(reader.status());
  EXPECT_EQ(2, reader.num_files());

  // We query some of the tensors
  {
    TensorShape shape;
    DataType type;
    EXPECT_TRUE(reader.HasTensor("test", &shape, &type));
    EXPECT_EQ("[4,5]", shape.DebugString());
    EXPECT_EQ(DataTypeToEnum<T>::v(), type);
    EXPECT_FALSE(reader.HasTensor("don't exist", nullptr, nullptr));
  }

  // Now we query some slices
  //
  // Slice #1 is an exact match
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("0,2:-");
    T expected[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
    U results[10];
    EXPECT_TRUE(reader.CopySliceData("test", s, results));
    for (int i = 0; i < 10; ++i) {
      EXPECT_EQ(expected[i], results[i]);
    }
  }

  // Slice #2 is a subset match
  //   .   .   .   .   .
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,1:-");
    T expected[] = {5, 6, 7, 8, 9};
    U results[5];
    EXPECT_TRUE(reader.CopySliceData("test", s, results));
    for (int i = 0; i < 5; ++i) {
      EXPECT_EQ(expected[i], results[i]);
    }
  }

  // Slice #4 includes the hole and so there is no match
  //   .   .   .   .   .
  //   .   .   7   8   9
  //   .   .  12  13  14
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,2:2,3");
    U results[6];
    EXPECT_FALSE(reader.CopySliceData("test", s, results));
  }
}

Project: tensorflow
File: tensor_slice_reader_test.cc

void CachedTensorSliceReaderTesterHelper(
    const TensorSliceWriter::CreateBuilderFunction& create_function,
    const TensorSliceReader::OpenTableFunction& open_function) {
  const string fname_base = io::JoinPath(testing::TmpDir(), "float_checkpoint");

  TensorShape shape({4, 5});

  // File #0 contains a slice that is the top two rows:
  //
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    const string fname = strings::StrCat(fname_base, "_0");
    TensorSliceWriter writer(fname, create_function);
    const float data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
    TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
    TF_CHECK_OK(writer.Add("test", shape, slice, data));
    TF_CHECK_OK(writer.Finish());
  }

  // File #1 contains two slices:
  //
  // slice #0 is the bottom left corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //  10  11  12   .   .
  //  15  16  17   .   .
  //
  // slice #1 is the bottom right corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .  18  19
  {
    const string fname = strings::StrCat(fname_base, "_1");
    TensorSliceWriter writer(fname, create_function);
    // slice #0
    {
      const float data[] = {10, 11, 12, 15, 16, 17};
      TensorSlice slice = TensorSlice::ParseOrDie("2,2:0,3");
      TF_CHECK_OK(writer.Add("test", shape, slice, data));
    }
    // slice #1
    {
      const float data[] = {18, 19};
      TensorSlice slice = TensorSlice::ParseOrDie("3,1:3,2");
      TF_CHECK_OK(writer.Add("test", shape, slice, data));
    }
    TF_CHECK_OK(writer.Finish());
  }

  // Notice that we leave a hole in the tensor
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   . (13) (14)
  //   .   .   .   .   .

  // Now we need to read the tensor slices
  TensorSliceReaderCache cache;
  const string filepattern = strings::StrCat(fname_base, "_*");
  const TensorSliceReader* reader = cache.GetReader(
      filepattern, open_function, TensorSliceReader::kLoadAllShards);
  EXPECT_TRUE(reader != nullptr);
  EXPECT_EQ(2, reader->num_files());

  // We query some of the tensors
  {
    TensorShape shape;
    DataType type;
    EXPECT_TRUE(reader->HasTensor("test", &shape, &type));
    EXPECT_EQ("[4,5]", shape.DebugString());
    EXPECT_EQ(DT_FLOAT, type);
    EXPECT_FALSE(reader->HasTensor("don't exist", nullptr, nullptr));
  }

  // Make sure the reader is cached.
  const TensorSliceReader* reader2 = cache.GetReader(
      filepattern, open_function, TensorSliceReader::kLoadAllShards);
  EXPECT_EQ(reader, reader2);

  reader = cache.GetReader("file_does_not_exist", open_function,
                           TensorSliceReader::kLoadAllShards);
  EXPECT_TRUE(reader == nullptr);
}

Project: tensorflow
File: tensor_slice.cc

void TensorSlice::UpdateToCover(const TensorSlice& other) {
  DCHECK_EQ(dims(), other.dims());
  for (int d = 0; d < dims(); ++d) {
    if (!IsFullAt(d)) {
      if (other.IsFullAt(d)) {
        starts_[d] = 0;
        lengths_[d] = kFullExtent;
      } else {
        const auto new_end = std::max(end(d), other.end(d));
        set_start(d, std::min(start(d), other.start(d)));
        set_length(d, new_end - start(d));
      }
    }
  }
}

Project: tensorflow
File: tensor_slice_reader.cc

void TensorSliceReader::LoadShard(int shard) const {
  CHECK_LT(shard, sss_.size());
  if (sss_[shard] || !status_.ok()) {
    return;  // Already loaded, or invalid.
  }
  string value;
  SavedTensorSlices sts;
  const string fname = fnames_[shard];
  VLOG(1) << "Reading meta data from file " << fname << "...";
  Table* table;
  Status s = open_function_(fname, &table);
  if (!s.ok()) {
    status_ = errors::DataLoss("Unable to open table file ", fname, ": ",
                               s.ToString());
    return;
  }
  sss_[shard].reset(table);
  if (!(table->Get(kSavedTensorSlicesKey, &value) &&
        ParseProtoUnlimited(&sts, value))) {
    status_ = errors::Internal(
        "Failed to find the saved tensor slices at the beginning of the "
        "checkpoint file: ",
        fname);
    return;
  }
  status_ = CheckVersions(sts.meta().versions(), TF_CHECKPOINT_VERSION,
                          TF_CHECKPOINT_VERSION_MIN_PRODUCER, "Checkpoint",
                          "checkpoint");
  if (!status_.ok()) return;
  for (const SavedSliceMeta& ssm : sts.meta().tensor()) {
    TensorShape ssm_shape;
    status_ = TensorShape::BuildTensorShapeBase(ssm.shape(), &ssm_shape);
    if (!status_.ok()) return;
    for (const TensorSliceProto& tsp : ssm.slice()) {
      TensorSlice ss_slice;
      status_ = TensorSlice::BuildTensorSlice(tsp, &ss_slice);
      if (!status_.ok()) return;
      status_ = RegisterTensorSlice(ssm.name(), ssm_shape, ssm.type(), fname,
                                    ss_slice, &tensors_);
      if (!status_.ok()) return;
    }
  }
}

Project: tensorflow
File: tensor_slice.cc

TensorSlice::TensorSlice(
    std::initializer_list<std::pair<int64_t, int64_t>> extents) {
  starts_.reserve(extents.size());
  lengths_.reserve(extents.size());
  for (const auto& e : extents) {
    starts_.push_back(e.first);
    lengths_.push_back(e.second);
  }
}

Project: tensorflow
File: tensor_slice.cc

void TensorSlice::ComputeRelative(const TensorSlice& sub,
                                  TensorSlice* relative) const {
  DCHECK_EQ(dims(), sub.dims());
  relative->SetFullSlice(dims());
  for (int d = 0; d < dims(); ++d) {
    if (IsFullAt(d)) {
      relative->set_start(d, sub.start(d));
      relative->set_length(d, sub.length(d));
    } else {
      // Otherwise the relative start is the difference between the start of
      // sub and the start of base
      relative->set_start(d, sub.start(d) - start(d));
      relative->set_length(d, sub.length(d));
    }
  }
}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, SliceTensorShape) {
  // A proper application
  {
    TensorSlice a = TensorSlice::ParseOrDie("1,1:-:4,1:2,6");
    TensorShape x({2, 4, 5, 8});
    TensorShape y;
    TF_EXPECT_OK(a.SliceTensorShape(x, &y));
    EXPECT_EQ("[1,4,1,6]", y.DebugString());
  }

  // An invalid application -- dimension 2 is out of bounds
  {
    TensorSlice a = TensorSlice::ParseOrDie("1,1:1,4:-:-");
    TensorShape x({2, 4, 5, 8});
    TensorShape y;
    Status s = a.SliceTensorShape(x, &y);
    EXPECT_EQ(s.code(), error::INTERNAL);
    EXPECT_TRUE(absl::StrContains(s.error_message(),
                                  "Extent in dimension 1 out of bounds: "
                                  "shape = [2,4,5,8], slice = 1,1:1,4:-:-"));
    EXPECT_EQ("[]", y.DebugString());
  }
}

Project: tensorflow
File: tensor_slice_reader.cc

void TensorSliceReader::LoadAllShards() const {
  VLOG(1) << "Loading all shards for " << filepattern_;
  for (size_t i = 0; i < fnames_.size() && status_.ok(); ++i) {
    LoadShard(i);
  }
  all_shards_loaded_ = true;
}

Project: tensorflow
File: tensor_slice.cc

int64_t TensorSlice::GetExtentLength(const TensorSliceProto::Extent& extent) {
  if (!HasExtentLength(extent)) return -1;
  return extent.length();
}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, ExtentLength) {
  TensorSliceProto proto;
  // Define ptxt outside ASSERT_TRUE call to work around bug in some
  // versions of gcc that breaks when you use raw string literals
  // inside macro expansions.
  //   See: https://gcc.gnu.org/bugzilla/show_bug.cgi?id=55971
  const char* ptxt = R"pb(
    extent {}
    extent { start: 0 length: 10 }
    extent { start: 14 length: 1 }
    extent {}
  )pb";
  ASSERT_TRUE(protobuf::TextFormat::ParseFromString(ptxt, &proto));
  EXPECT_FALSE(TensorSlice::HasExtentLength(proto.extent(0)));
  EXPECT_TRUE(TensorSlice::HasExtentLength(proto.extent(1)));
  EXPECT_TRUE(TensorSlice::HasExtentLength(proto.extent(2)));
  EXPECT_FALSE(TensorSlice::HasExtentLength(proto.extent(3)));
  EXPECT_EQ(-1, TensorSlice::GetExtentLength(proto.extent(0)));
  EXPECT_EQ(10, TensorSlice::GetExtentLength(proto.extent(1)));
  EXPECT_EQ(1, TensorSlice::GetExtentLength(proto.extent(2)));
  EXPECT_EQ(-1, TensorSlice::GetExtentLength(proto.extent(3)));
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(TensorSliceReaderTest, UnsupportedTensorType) {
  const string fname = io::JoinPath(testing::TmpDir(), "int32_ref_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TensorShape shape({4, 5});
  TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
  TF_CHECK_OK(writer.Add("test", shape, slice, data));
  TF_CHECK_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
    if (sts.has_meta()) {
      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
        tensor.set_type(DT_INT32_REF);
      }
    }
    return sts.SerializeAsString();
  });

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  TF_CHECK_OK(reader.status());

  // The tensor should be present, but loading it should fail due to the
  // unsupported type.
  EXPECT_TRUE(reader.HasTensor("test", nullptr, nullptr));
  std::unique_ptr<Tensor> tensor;
  EXPECT_FALSE(reader.GetTensor("test", &tensor).ok());
}

Project: tensorflow
File: tensor_slice_reader.cc

const TensorSliceSet* TensorSliceReader::FindTensorSlice(
    const string& name, const TensorSlice& slice,
    std::vector<std::pair<TensorSlice, string>>* details) const {
  const TensorSliceSet* tss = gtl::FindPtrOrNull(tensors_, name);
  if (tss && !tss->QueryMeta(slice, details)) {
    return nullptr;
  }
  return tss;
}

Project: tensorflow
File: tensor_slice_reader.cc

  ~TensorSliceReaderTable() override {
    delete table_;
    delete file_;
  }

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, Intersection) {
  // "EVERYTHING" intersects with everything
  {
    TensorSlice a = TensorSlice::ParseOrDie("-:-");
    TensorSlice b = TensorSlice::ParseOrDie("1,2:3,4");
    TensorSlice c;
    EXPECT_TRUE(a.Intersect(b, &c));
    EXPECT_EQ("1,2:3,4", c.DebugString());
  }

  {
    TensorSlice a = TensorSlice::ParseOrDie("-:-");
    TensorSlice b = TensorSlice::ParseOrDie("1,2:3,4");
    TensorSlice c;
    EXPECT_TRUE(b.Intersect(a, &c));
    EXPECT_EQ("1,2:3,4", c.DebugString());
  }

  // Overlap at all dimensions
  {
    TensorSlice a = TensorSlice::ParseOrDie("1,5:2,6:3,7:5,10");
    TensorSlice b = TensorSlice::ParseOrDie("1,2:3,4:9,10:12,1");
    TensorSlice c;
    EXPECT_TRUE(a.Intersect(b, &c));
    EXPECT_EQ("1,2:3,4:9,1:12,1", c.DebugString());
  }

  // A mixture of everything and non-trivial slices
  {
    TensorSlice a = TensorSlice::ParseOrDie("-:1,1");
    TensorSlice b = TensorSlice::ParseOrDie("-:0,2");
    TensorSlice c;
    EXPECT_TRUE(a.Intersect(b, &c));
    EXPECT_EQ("-:1,1", c.DebugString());
  }

  // No overlap on dimension 3: "3,1" and "4,5" don't intersect
  {
    TensorSlice a = TensorSlice::ParseOrDie("1,2:3,1:5,6");
    TensorSlice b = TensorSlice::ParseOrDie("1,3:4,5:1,6");
    TensorSlice c;
    EXPECT_FALSE(a.Intersect(b, &c));
    EXPECT_EQ("", c.DebugString());
  }
  // No intersection when there are different dimensions
  {
    TensorSlice a = TensorSlice::ParseOrDie("1,2:3,1:-");
    TensorSlice b = TensorSlice::ParseOrDie("-:-");
    TensorSlice c;
    EXPECT_FALSE(a.Intersect(b, &c));
    EXPECT_EQ("", c.DebugString());
  }
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(CheckpointVersionTest, MinProducer) {
  VersionDef versions;
  versions.set_producer(TF_CHECKPOINT_VERSION_MIN_PRODUCER - 1);
  VersionTest(versions, strings::StrCat("Checkpoint producer version ",
                                        TF_CHECKPOINT_VERSION_MIN_PRODUCER - 1,
                                        " below min producer ",
                                        TF_CHECKPOINT_VERSION_MIN_PRODUCER,
                                        " supported by TensorFlow"));
}

Project: tensorflow
File: tensor_slice.cc

string TensorSlice::DebugString() const {
  string buffer;
  bool first = true;
  for (int d = 0; d < dims(); ++d) {
    if (!first) {
      buffer.append(":");
    }
    if (IsFullAt(d)) {
      buffer.append("-");
    } else {
      strings::StrAppend(&buffer, starts_[d], ",", lengths_[d]);
    }
    first = false;
  }
  return buffer;
}

Project: tensorflow
File: tensor_slice.cc

bool TensorSlice::Intersect(const TensorSlice& other,
                            TensorSlice* result) const {
  // First, if two slices have different ranks, they obviously don't overlap
  // -- in fact they are not compatible.
  if (dims() != other.dims()) {
    return false;
  }

  // Setting the result to the right dimension
  if (result) {
    result->SetFullSlice(dims());
  }
  // The two slices overlap if they overlap in all dimensions.
  for (int d = 0; d < dims(); ++d) {
    if (IsFullAt(d)) {
      if (result) {
        result->set_start(d, other.start(d));
        result->set_length(d, other.length(d));
      }
    } else if (other.IsFullAt(d)) {
      if (result) {
        result->set_start(d, start(d));
        result->set_length(d, length(d));
      }
    } else {
      // If we have an intersection here, it should have a start that is the
      // max of the two starts and an end that is the min of the two ends.
      int64_t s = std::max(start(d), other.start(d));
      int64_t l = std::min(end(d), other.end(d)) - s;
      if (l > 0) {
        // We have a real intersection
        if (result) {
          result->set_start(d, s);
          result->set_length(d, l);
        }
      } else {
        // We don't have an intersection for this dimension -- thus we don't
        // have any intersection at all.
        if (result) {
          result->Clear();
        }
        return false;
      }
    }
  }
  // If we are here, we know there is overlap in every dimension.
  return true;
}

Project: tensorflow
File: tensor_slice.cc

void TensorSlice::Extend(int dim) {
  int old_dim = dims();
  DCHECK_LE(old_dim, dim);
  starts_.resize(dim);
  lengths_.resize(dim);
  for (int d = old_dim; d < dim; ++d) {
    starts_[d] = 0;
    lengths_[d] = kFullExtent;
  }
}

Project: tensorflow
File: tensor_slice_test.cc

TEST(TensorSliceTest, Basic) {
  {
    // Repeatedly setting FullSlice should work.
    TensorSlice s(3);
    EXPECT_EQ("-:-:-", s.DebugString());
    EXPECT_TRUE(s.IsFull());

    s.SetFullSlice(4);
    EXPECT_EQ("-:-:-:-", s.DebugString());
    EXPECT_TRUE(s.IsFull());
  }
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(CheckpointVersionTest, BadConsumer) {
  VersionDef versions;
  versions.set_producer(TF_CHECKPOINT_VERSION + 1);
  versions.add_bad_consumers(TF_CHECKPOINT_VERSION);
  VersionTest(
      versions,
      strings::StrCat(
          "Checkpoint disallows consumer version ", TF_CHECKPOINT_VERSION,
          ".  Please upgrade TensorFlow: this version is likely buggy."));
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(CachedTensorSliceReaderTest, SimpleFloat) {
  CachedTensorSliceReaderTesterHelper(CreateTableTensorSliceBuilder,
                                      OpenTableTensorSliceReader);
}

Project: tensorflow
File: tensor_slice_reader_test.cc

TEST(TensorSliceReaderTest, MissingTensorType) {
  const string fname = io::JoinPath(testing::TmpDir(), "invalid_checkpoint");
  TensorSliceWriter writer(fname, CreateTableTensorSliceBuilder);
  const int32 data[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TensorShape shape({4, 5});
  TensorSlice slice = TensorSlice::ParseOrDie("0,2:-");
  TF_CHECK_OK(writer.Add("test", shape, slice, data));
  TF_CHECK_OK(writer.Finish());

  MutateSavedTensorSlices(fname, [](SavedTensorSlices sts) {
    if (sts.has_meta()) {
      for (auto& tensor : *sts.mutable_meta()->mutable_tensor()) {
        tensor.clear_type();
      }
    }
    return sts.SerializeAsString();
  });

  TensorSliceReader reader(fname, OpenTableTensorSliceReader);
  TF_CHECK_OK(reader.status());

  // The tensor should be present, but loading it should fail due to the
  // unset (invalid) type.
  EXPECT_TRUE(reader.HasTensor("test", nullptr, nullptr));
  std::unique_ptr<Tensor> tensor;
  EXPECT_FALSE(reader.GetTensor("test", &tensor).ok());
}

Project: tensorflow
File: tensor_slice_reader.cc

TensorSliceReader::TensorSliceReader(const string& filepattern,
                                     OpenTableFunction open_function,
                                     int preferred_shard)
    : filepattern_(filepattern), open_function_(std::move(open_function)) {
  VLOG(1) << "TensorSliceReader for " << filepattern;
  Status s = Env::Default()->GetMatchingPaths(filepattern, &fnames_);
  if (!s.ok()) {
    status_ = errors::InvalidArgument(
        "Unsuccessful TensorSliceReader constructor: "
        "Failed to get matching files on ",
        filepattern, ": ", s.ToString());
    return;
  }
  if (fnames_.empty()) {
    status_ = errors::NotFound(
        "Unsuccessful TensorSliceReader constructor: "
        "Failed to find any matching files for ",
        filepattern);
    return;
  }
  sss_.resize(fnames_.size());
  for (size_t shard = 0; shard < fnames_.size(); ++shard) {
    fname_to_index_.insert(std::make_pair(fnames_[shard], shard));
  }
  if (preferred_shard == kLoadAllShards || fnames_.size() == 1 ||
      static_cast<size_t>(preferred_shard) >= fnames_.size()) {
    LoadAllShards();
  } else {
    VLOG(1) << "Loading shard " << preferred_shard << " for " << filepattern_;
    LoadShard(preferred_shard);
  }
}