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
58
59
60
61
62
63
64
65
66
|
/*
* This file is part of the micropython-ulab project,
*
* https://github.com/v923z/micropython-ulab
*
* The MIT License (MIT)
*
* Copyright (c) 2021 Zoltán Vörös
*
*/
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "py/obj.h"
#include "py/runtime.h"
#include "ndarray_iter.h"
#ifdef NDARRAY_HAS_FLATITER
mp_obj_t ndarray_flatiter_make_new(mp_obj_t self_in) {
ndarray_flatiter_t *flatiter = m_new_obj(ndarray_flatiter_t);
flatiter->base.type = &ndarray_flatiter_type;
flatiter->iternext = ndarray_flatiter_next;
flatiter->ndarray = MP_OBJ_TO_PTR(self_in);
flatiter->cur = 0;
return flatiter;
}
mp_obj_t ndarray_flatiter_next(mp_obj_t self_in) {
ndarray_flatiter_t *self = MP_OBJ_TO_PTR(self_in);
ndarray_obj_t *ndarray = MP_OBJ_TO_PTR(self->ndarray);
uint8_t *array = (uint8_t *)ndarray->array;
if(self->cur < ndarray->len) {
uint32_t remainder = self->cur;
uint8_t i = ULAB_MAX_DIMS - 1;
do {
size_t div = (remainder / ndarray->shape[i]);
array += remainder * ndarray->strides[i];
remainder -= div * ndarray->shape[i];
i--;
} while(i > ULAB_MAX_DIMS - ndarray->ndim);
self->cur++;
return ndarray_get_item(ndarray, array);
}
return MP_OBJ_STOP_ITERATION;
}
mp_obj_t ndarray_new_flatiterator(mp_obj_t flatiter_in, mp_obj_iter_buf_t *iter_buf) {
assert(sizeof(ndarray_flatiter_t) <= sizeof(mp_obj_iter_buf_t));
ndarray_flatiter_t *iter = (ndarray_flatiter_t *)iter_buf;
ndarray_flatiter_t *flatiter = MP_OBJ_TO_PTR(flatiter_in);
iter->base.type = &mp_type_polymorph_iter;
iter->iternext = ndarray_flatiter_next;
iter->ndarray = flatiter->ndarray;
iter->cur = 0;
return MP_OBJ_FROM_PTR(iter);
}
mp_obj_t ndarray_get_flatiterator(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) {
return ndarray_new_flatiterator(o_in, iter_buf);
}
#endif /* NDARRAY_HAS_FLATITER */
|
