Skip to content

Commit d62a934

Browse files
kirklandsignfacebook-github-bot
authored andcommitted
reduce_util: clamp the apply fn start/end
Summary: In helper apply_over_dim_list, we can clamp the start/end, and skip if we don't need to apply fn to any elements. Reviewed By: manuelcandales Differential Revision: D48211547 fbshipit-source-id: 311b271e15baf46fbccaa4bacac072981abe6f61
1 parent 3a92785 commit d62a934

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

kernels/portable/cpu/util/reduce_util.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -299,10 +299,10 @@ void apply_over_dim_list(
299299
}
300300

301301
const size_t iter_length = get_reduced_dim_product(in, dim_list);
302-
ET_CHECK_VALID_IX(start, iter_length);
303-
ET_CHECK_VALID_IX(end, iter_length);
304-
const size_t ustart = ET_NORMALIZE_IX(start, iter_length);
305-
const size_t uend = ET_NORMALIZE_IX(end, iter_length);
302+
const size_t normalized_start = ET_NORMALIZE_IX(start, iter_length);
303+
const size_t normalized_end = ET_NORMALIZE_IX(end, iter_length);
304+
const size_t ustart = std::max(normalized_start, size_t(0));
305+
const size_t uend = std::min(normalized_end, iter_length - 1);
306306

307307
// If dim_list is null or empty, iterate over the entire tensor
308308
if (!dim_list.has_value() || dim_list.value().size() == 0) {

0 commit comments

Comments
 (0)