Skip to content
Permalink
Browse files Browse the repository at this point in the history
[lite] Add some safety checks to avoid out of bound access for sparsi…
…ty format

PiperOrigin-RevId: 416910386
Change-Id: Ic0b4dc048dc4b5a6309c572b8c4c9f776e4db60a
  • Loading branch information
karimnosseir authored and tensorflower-gardener committed Dec 16, 2021
1 parent 3e49ff6 commit 6364463
Showing 1 changed file with 11 additions and 7 deletions.
Expand Up @@ -282,10 +282,12 @@ void FormatConverter<T>::InitSparseToDenseConverter(
block_size_.resize(block_map_.size());
for (int i = 0; i < original_rank; i++) {
if (block_dim < block_map_.size() && block_map_[block_dim] == i) {
int orig_dim = traversal_order_[original_rank + block_dim];
block_size_[block_dim] = dense_size[orig_dim];
blocked_shape_[i] = dense_shape_[i] / dense_size[orig_dim];
block_dim++;
if (original_rank + block_dim < traversal_order_.size()) {
int orig_dim = traversal_order_[original_rank + block_dim];
block_size_[block_dim] = dense_size[orig_dim];
blocked_shape_[i] = dense_shape_[i] / dense_size[orig_dim];
block_dim++;
}
} else {
blocked_shape_[i] = dense_shape_[i];
}
Expand Down Expand Up @@ -328,13 +330,15 @@ void FormatConverter<T>::Populate(const T* src_data, std::vector<int> indices,
Populate(src_data, indices, level + 1, prev_idx * shape_of_level + i,
src_data_ptr, dest_data);
}
} else {
} else if (prev_idx + 1 < dim_metadata_[metadata_idx].size()) {
const auto& array_segments = dim_metadata_[metadata_idx];
const auto& array_indices = dim_metadata_[metadata_idx + 1];
for (int i = array_segments[prev_idx]; i < array_segments[prev_idx + 1];
i++) {
indices[level] = array_indices[i];
Populate(src_data, indices, level + 1, i, src_data_ptr, dest_data);
if (i < array_indices.size() && level < indices.size()) {
indices[level] = array_indices[i];
Populate(src_data, indices, level + 1, i, src_data_ptr, dest_data);
}
}
}
}
Expand Down

0 comments on commit 6364463

Please sign in to comment.