core: Fix IRRT implementation of ndarray_flatten_index

This commit is contained in:
David Mak 2024-02-15 14:57:41 +08:00
parent 282a3e1911
commit 8c7e44098a
1 changed files with 10 additions and 8 deletions

View File

@ -278,12 +278,13 @@ uint32_t __nac3_ndarray_flatten_index(
) { ) {
uint32_t idx = 0; uint32_t idx = 0;
uint32_t stride = 1; uint32_t stride = 1;
for (uint32_t i = num_dims - 1; i-- >= 0; ) { for (uint32_t i = 0; i < num_dims; ++i) {
if (i < num_indices) { uint32_t ri = num_dims - i - 1;
idx += (stride * indices[i]); if (ri < num_indices) {
idx += (stride * indices[ri]);
} }
stride *= dims[i]; stride *= dims[ri];
} }
return idx; return idx;
} }
@ -296,12 +297,13 @@ uint64_t __nac3_ndarray_flatten_index64(
) { ) {
uint64_t idx = 0; uint64_t idx = 0;
uint64_t stride = 1; uint64_t stride = 1;
for (uint64_t i = num_dims - 1; i-- >= 0; ) { for (uint64_t i = 0; i < num_dims; ++i) {
if (i < num_indices) { uint64_t ri = num_dims - i - 1;
idx += (stride * indices[i]); if (ri < num_indices) {
idx += (stride * indices[ri]);
} }
stride *= dims[i]; stride *= dims[ri];
} }
return idx; return idx;
} }