MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / deduce_layout

Function deduce_layout

imperative/src/impl/ops/subtensor.cpp:85–143  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

83};
84
85TensorLayout deduce_layout(
86 TensorLayout src, std::vector<std::tuple<int8_t, bool, bool, bool, bool>> items,
87 std::vector<std::tuple<int32_t, int32_t, int32_t, int32_t>> slice_items) {
88 auto tostr = [](int v) -> std::string { return std::to_string(v); };
89
90 for (int i = items.size() - 1; i >= 0; i--) {
91 auto&& [axis, b_flag, e_flag, s_flag, idx_flag] = items[i];
92 auto&& [b_val, e_val, s_val, ax_val] = slice_items[i];
93 int shape_axis = src.shape[axis];
94 int slice_step = s_val == INT_MAX ? 1 : s_val;
95 int slice_start = b_val == INT_MIN ? 0 : b_val;
96 int slice_stop = e_val == INT_MAX ? shape_axis : e_val;
97 if (slice_step > 0) {
98 slice_start = mod_size(slice_start, shape_axis);
99 slice_stop = mod_size(slice_stop, shape_axis);
100 slice_stop = std::min(slice_stop, shape_axis);
101 slice_start = std::min(slice_start, slice_stop);
102 mgb_assert(
103 (slice_start >= 0 && slice_stop >= slice_start &&
104 slice_stop <= shape_axis),
105 "index out of bound: layout=%s; request begin=%s end=%s step=%s "
106 "axis=%s",
107 src.to_string().c_str(), tostr(slice_start).c_str(),
108 tostr(slice_stop).c_str(), tostr(slice_step).c_str(),
109 tostr(axis).c_str());
110 } else {
111 slice_start = b_val == INT_MIN ? shape_axis - 1 : b_val;
112 slice_start = mod_size(slice_start, shape_axis);
113 slice_stop = e_val == INT_MAX ? -1 : mod_size(e_val, shape_axis);
114 slice_start = std::min(slice_start, std::max(shape_axis - 1, 0));
115 slice_stop = std::min(slice_stop, slice_start);
116 mgb_assert(
117 (slice_step < 0 && slice_start >= 0 && slice_stop <= slice_start &&
118 slice_start < shape_axis && slice_stop >= -1),
119 "index out of bound: layout=%s; request begin=%s end=%s step=%s "
120 "axis=%s",
121 src.to_string().c_str(), tostr(slice_start).c_str(),
122 tostr(slice_stop).c_str(), tostr(slice_step).c_str(),
123 tostr(axis).c_str());
124 }
125 int abs_step = std::abs(slice_step);
126 if (axis < 0) {
127 axis = axis + src.ndim;
128 };
129
130 if (idx_flag == true) {
131 if (src.ndim == 1) {
132 src.shape[0] = 1;
133 } else {
134 src.remove_axis_inplace(axis);
135 }
136 } else {
137 src.shape[axis] =
138 (std::abs(slice_stop - slice_start) + abs_step - 1) / abs_step;
139 src.stride[axis] *= slice_step;
140 }
141 }
142 return src;

Callers 12

apply_on_physical_tensorFunction · 0.85
execMethod · 0.85
check_execMethod · 0.85
check_execMethod · 0.85
check_execMethod · 0.85
check_execMethod · 0.85
check_execMethod · 0.85
check_execMethod · 0.85
check_execMethod · 0.85
check_layoutMethod · 0.85

Calls 7

to_stringFunction · 0.85
maxFunction · 0.85
remove_axis_inplaceMethod · 0.80
minFunction · 0.50
absFunction · 0.50
sizeMethod · 0.45
to_stringMethod · 0.45

Tested by

no test coverage detected