1 | #include <c10/core/WrapDimMinimal.h> |
2 | |
3 | namespace c10 { |
4 | namespace detail { |
5 | |
6 | template <typename T> |
7 | T maybe_wrap_dim_slow(T dim, T dim_post_expr, bool wrap_scalar) { |
8 | TORCH_CHECK_INDEX( |
9 | dim_post_expr >= 0, "Rank cannot be negative but got " , dim_post_expr); |
10 | |
11 | if (dim_post_expr == 0) { |
12 | TORCH_CHECK_INDEX( |
13 | wrap_scalar, |
14 | "Dimension specified as " , |
15 | dim, |
16 | " but tensor has no dimensions" ); |
17 | return c10::maybe_wrap_dim( |
18 | std::move(dim), /*dim_post_expr=*/1, /*wrap_scalar=*/false); |
19 | } |
20 | |
21 | T min = dim_post_expr * -1; |
22 | T max = dim_post_expr - 1; |
23 | TORCH_CHECK_INDEX( |
24 | min <= dim && dim <= max, |
25 | "Dimension out of range (expected to be in range of [" , |
26 | min, |
27 | ", " , |
28 | max, |
29 | "], but got " , |
30 | dim, |
31 | ")" ); |
32 | |
33 | TORCH_INTERNAL_ASSERT( |
34 | false, "should never reach here as dim should be out-of-bounds" ); |
35 | } |
36 | |
37 | // Explicitly instantiate the template at the two types it will be used |
38 | template C10_API int64_t |
39 | maybe_wrap_dim_slow(int64_t dim, int64_t dim_post_expr, bool wrap_scalar); |
40 | template C10_API SymInt |
41 | maybe_wrap_dim_slow(SymInt dim, SymInt dim_post_expr, bool wrap_scalar); |
42 | |
43 | } // namespace detail |
44 | } // namespace c10 |
45 | |