1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
|
#include <torch/csrc/jit/runtime/slice_indices_adjust.h>
#include <c10/util/Exception.h>
#include <cstdint>
namespace torch {
namespace jit {
int64_t slice_indices_adjust(
int64_t length,
int64_t* start,
int64_t* stop,
int64_t step) {
TORCH_CHECK(step != 0, "List slice should have non-zero step")
TORCH_CHECK(step >= -INT64_MAX, "List slice step is out of bounds")
// Comes from PySlice_Unpack.
if (*start == INT64_MAX) {
*start = (step < 0) ? INT64_MAX : 0;
}
if (*stop == INT64_MAX) {
*stop = (step < 0) ? INT64_MIN : INT64_MAX;
}
// Comes from PySlice_AdjustIndices.
if (*start < 0) {
*start += length;
if (*start < 0) {
*start = (step < 0) ? -1 : 0;
}
} else if (*start >= length) {
*start = (step < 0) ? length - 1 : length;
}
if (*stop < 0) {
*stop += length;
if (*stop < 0) {
*stop = (step < 0) ? -1 : 0;
}
} else if (*stop >= length) {
*stop = (step < 0) ? length - 1 : length;
}
if (step < 0) {
if (*stop < *start) {
return (*start - *stop - 1) / (-step) + 1;
}
} else {
if (*start < *stop) {
return (*stop - *start - 1) / step + 1;
}
}
return 0;
}
} // namespace jit
} // namespace torch
|
