diff options
-rw-r--r-- | ndarray.c | 12 | ||||
-rw-r--r-- | ndarray.h | 2 |
2 files changed, 8 insertions, 6 deletions
@@ -128,8 +128,8 @@ int mg2di_ndarray_slice(NDArray **result, NDArray *src, const Slice *s = &slice[i]; ptrdiff_t start, end; - if (s->step != 1) { - ret = -ENOSYS; + if (s->step < 1) { + ret = s->step ? -ENOSYS : -EINVAL; goto fail; } if (s->start == PTRDIFF_MAX) { @@ -148,10 +148,12 @@ int mg2di_ndarray_slice(NDArray **result, NDArray *src, ret = -EINVAL; goto fail; } - - offset += src->stride[i] * start; - a->priv->shape[i] = end - start; } + + offset += src->stride[i] * start; + + a->priv->stride[i] *= s->step; + a->priv->shape[i] = (end - start) / s->step; } a->dims = src->dims; @@ -25,7 +25,7 @@ typedef struct Slice { ptrdiff_t start; ptrdiff_t end; - size_t step; + ptrdiff_t step; } Slice; #define SLICE(start, end, step) ((Slice){ start, end, step }) |