| 27 | |
| 28 | @numba.jit(nopython=nopython, nogil=nogil, parallel=parallel) |
| 29 | def nb_looper(values, axis, *args): |
| 30 | # Operate on the first row/col in order to get |
| 31 | # the output shape |
| 32 | if axis == 0: |
| 33 | first_elem = values[:, 0] |
| 34 | dim0 = values.shape[1] |
| 35 | else: |
| 36 | first_elem = values[0] |
| 37 | dim0 = values.shape[0] |
| 38 | res0 = nb_compat_func(first_elem, *args) |
| 39 | # Use np.asarray to get shape for |
| 40 | # https://github.com/numba/numba/issues/4202#issuecomment-1185981507 |
| 41 | # Use tuple concatenation; numba doesn't support tuple unpacking syntax |
| 42 | buf_shape = (dim0,) + np.atleast_1d(np.asarray(res0)).shape # noqa: RUF005 |
| 43 | if axis == 0: |
| 44 | buf_shape = buf_shape[::-1] |
| 45 | buff = np.empty(buf_shape) |
| 46 | |
| 47 | if axis == 1: |
| 48 | buff[0] = res0 |
| 49 | for i in numba.prange(1, values.shape[0]): |
| 50 | buff[i] = nb_compat_func(values[i], *args) |
| 51 | else: |
| 52 | buff[:, 0] = res0 |
| 53 | for j in numba.prange(1, values.shape[1]): |
| 54 | buff[:, j] = nb_compat_func(values[:, j], *args) |
| 55 | return buff |
| 56 | |
| 57 | return nb_looper |
| 58 | |