Skip to content

Commit 6364463

Browse files
karimnosseirtensorflower-gardener
authored andcommitted
[lite] Add some safety checks to avoid out of bound access for sparsity format
PiperOrigin-RevId: 416910386 Change-Id: Ic0b4dc048dc4b5a6309c572b8c4c9f776e4db60a
1 parent 3e49ff6 commit 6364463

File tree

1 file changed

+11
-7
lines changed

1 file changed

+11
-7
lines changed

Diff for: tensorflow/lite/kernels/internal/utils/sparsity_format_converter.cc

+11-7
Original file line numberDiff line numberDiff line change
@@ -282,10 +282,12 @@ void FormatConverter<T>::InitSparseToDenseConverter(
282282
block_size_.resize(block_map_.size());
283283
for (int i = 0; i < original_rank; i++) {
284284
if (block_dim < block_map_.size() && block_map_[block_dim] == i) {
285-
int orig_dim = traversal_order_[original_rank + block_dim];
286-
block_size_[block_dim] = dense_size[orig_dim];
287-
blocked_shape_[i] = dense_shape_[i] / dense_size[orig_dim];
288-
block_dim++;
285+
if (original_rank + block_dim < traversal_order_.size()) {
286+
int orig_dim = traversal_order_[original_rank + block_dim];
287+
block_size_[block_dim] = dense_size[orig_dim];
288+
blocked_shape_[i] = dense_shape_[i] / dense_size[orig_dim];
289+
block_dim++;
290+
}
289291
} else {
290292
blocked_shape_[i] = dense_shape_[i];
291293
}
@@ -328,13 +330,15 @@ void FormatConverter<T>::Populate(const T* src_data, std::vector<int> indices,
328330
Populate(src_data, indices, level + 1, prev_idx * shape_of_level + i,
329331
src_data_ptr, dest_data);
330332
}
331-
} else {
333+
} else if (prev_idx + 1 < dim_metadata_[metadata_idx].size()) {
332334
const auto& array_segments = dim_metadata_[metadata_idx];
333335
const auto& array_indices = dim_metadata_[metadata_idx + 1];
334336
for (int i = array_segments[prev_idx]; i < array_segments[prev_idx + 1];
335337
i++) {
336-
indices[level] = array_indices[i];
337-
Populate(src_data, indices, level + 1, i, src_data_ptr, dest_data);
338+
if (i < array_indices.size() && level < indices.size()) {
339+
indices[level] = array_indices[i];
340+
Populate(src_data, indices, level + 1, i, src_data_ptr, dest_data);
341+
}
338342
}
339343
}
340344
}

0 commit comments

Comments
 (0)