From 556c53a155b40672868141431382027b2123fefe Mon Sep 17 00:00:00 2001 From: NAITOH Jun Date: Sun, 1 Sep 2019 05:51:56 +0900 Subject: [PATCH] at() method was rewritten in C. --- ext/cumo/include/cumo/intern.h | 1 + ext/cumo/narray/gen/spec.rb | 1 + ext/cumo/narray/gen/tmpl/at.c | 34 +++++ ext/cumo/narray/index.c | 221 +++++++++++++++++++++++++++++++- ext/cumo/narray/index_kernel.cu | 84 ++++++++++++ lib/cumo/narray/extra.rb | 50 -------- test/narray_test.rb | 35 +++++ 7 files changed, 369 insertions(+), 57 deletions(-) create mode 100644 ext/cumo/narray/gen/tmpl/at.c diff --git a/ext/cumo/include/cumo/intern.h b/ext/cumo/include/cumo/intern.h index 82cbef97..fea02549 100644 --- a/ext/cumo/include/cumo/intern.h +++ b/ext/cumo/include/cumo/intern.h @@ -79,6 +79,7 @@ void cumo_na_parse_enumerator_step(VALUE enum_obj, VALUE *pstep); // used in aref, aset int cumo_na_get_result_dimension(VALUE self, int argc, VALUE *argv, ssize_t stride, size_t *pos_idx); VALUE cumo_na_aref_main(int nidx, VALUE *idx, VALUE self, int keep_dim, int result_nd, size_t pos); +VALUE cumo_na_at_main(int nidx, VALUE *idx, VALUE self, int keep_dim, int result_nd, size_t pos); // defined in array, used in math VALUE cumo_na_ary_composition_dtype(VALUE ary); diff --git a/ext/cumo/narray/gen/spec.rb b/ext/cumo/narray/gen/spec.rb index 4cfc071a..c7d40e83 100644 --- a/ext/cumo/narray/gen/spec.rb +++ b/ext/cumo/narray/gen/spec.rb @@ -144,6 +144,7 @@ def_method "aref", op:"[]" def_method "aref_cpu" def_method "aset", op:"[]=" +def_method "at" def_method "coerce_cast" def_method "to_a" diff --git a/ext/cumo/narray/gen/tmpl/at.c b/ext/cumo/narray/gen/tmpl/at.c new file mode 100644 index 00000000..49c06d18 --- /dev/null +++ b/ext/cumo/narray/gen/tmpl/at.c @@ -0,0 +1,34 @@ +/* + Multi-dimensional array indexing. + Same as [] for one-dimensional NArray. + Similar to numpy's tuple indexing, i.e., `a[[1,2,..],[3,4,..]]` + @overload at(*indices) + @param [Numeric,Range,etc] *indices Multi-dimensional Index Arrays. + @return [Cumo::NArray::<%=class_name%>] one-dimensional NArray view. + + @example + x = Cumo::DFloat.new(3,3,3).seq + => Cumo::DFloat#shape=[3,3,3] + [[[0, 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]]] + + x.at([0,1,2],[0,1,2],[-1,-2,-3]) + => Cumo::DFloat(view)#shape=[3] + [2, 13, 24] + */ +static VALUE +<%=c_func(-1)%>(int argc, VALUE *argv, VALUE self) +{ + int result_nd; + size_t pos; + + result_nd = cumo_na_get_result_dimension(self, argc, argv, sizeof(dtype), &pos); + return cumo_na_at_main(argc, argv, self, 0, result_nd, pos); +} diff --git a/ext/cumo/narray/index.c b/ext/cumo/narray/index.c index cd0f959e..4fb53580 100644 --- a/ext/cumo/narray/index.c +++ b/ext/cumo/narray/index.c @@ -568,6 +568,187 @@ cumo_na_index_aref_naview(cumo_narray_view_t *na1, cumo_narray_view_t *na2, na2->base.size = total; } +void cumo_na_index_at_nadata_index_stride_add_kernel_launch(size_t *idx, size_t *idx1, ssize_t s1, uint64_t n); +void cumo_na_index_at_nadata_index_beg_step_stride_kernel_launch(size_t *idx, size_t beg, ssize_t step, ssize_t s1, uint64_t n); +void cumo_na_index_at_nadata_index_beg_step_stride_add_kernel_launch(size_t *idx, size_t beg, ssize_t step, ssize_t s1, uint64_t n); + +static void +cumo_na_index_at_nadata(cumo_narray_data_t *na1, cumo_narray_view_t *na2, + cumo_na_index_arg_t *q, ssize_t elmsz, int ndim, int keep_dim) +{ + int i; + ssize_t size = q[ndim-1].n; + ssize_t stride1; + ssize_t *strides_na1; + size_t *index; + ssize_t beg, step; + int use_cumo_cuda_runtime_malloc = 0; + + strides_na1 = ALLOCA_N(ssize_t, na1->base.ndim); + cumo_na_get_strides_nadata(na1, strides_na1, elmsz); + + if (q[ndim-1].idx != NULL) { + index = q[ndim-1].idx; + } else { + //index = ALLOC_N(size_t, size); + index = (size_t*)cumo_cuda_runtime_malloc(sizeof(size_t)*size); + use_cumo_cuda_runtime_malloc = 1; + } + CUMO_SDX_SET_INDEX(na2->stridx[0], index); + + for (i=ndim-1; i>=0; i--) { + stride1 = strides_na1[q[i].orig_dim]; + if (i==ndim-1) { + if (size == 0) { + rb_raise(cumo_na_eShapeError, "cannot get element of empty array"); + } + } else { + if (size != q[i].n) { + rb_raise(cumo_na_eShapeError, "index array sizes mismatch"); + } + } + + if (q[i].idx != NULL) { + if (i==ndim-1) { + cumo_na_index_aref_nadata_index_stride_kernel_launch(index, stride1, size); + } else { + cumo_na_index_at_nadata_index_stride_add_kernel_launch(index, q[i].idx, stride1, size); + } + q[i].idx = NULL; + } else { + beg = q[i].beg; + step = q[i].step; + if (i==ndim-1) { + cumo_na_index_at_nadata_index_beg_step_stride_kernel_launch(index, beg, step, stride1, size); + } else { + cumo_na_index_at_nadata_index_beg_step_stride_add_kernel_launch(index, beg, step, stride1, size); + } + } + + } + na2->base.size = size; + na2->base.shape[0] = size; + if (use_cumo_cuda_runtime_malloc) { + CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("index", "cumo_na_index_at_nadata"); + cumo_cuda_runtime_check_status(cudaDeviceSynchronize()); + } +} + +void cumo_na_index_at_naview_index_index_index_add_kernel_launch(size_t *idx, size_t *idx1, size_t *idx2, uint64_t n); +void cumo_na_index_at_naview_index_index_beg_step_add_kernel_launch(size_t *idx, size_t *idx1, size_t beg, ssize_t step, uint64_t n); +void cumo_na_index_at_naview_index_stride_last_add_kernel_launch(size_t *idx, ssize_t s1, size_t last, uint64_t n); + +static void +cumo_na_index_at_naview(cumo_narray_view_t *na1, cumo_narray_view_t *na2, + cumo_na_index_arg_t *q, ssize_t elmsz, int ndim, int keep_dim) +{ + int i; + size_t *index; + ssize_t size = q[ndim-1].n; + int use_cumo_cuda_runtime_malloc = 0; + + if (q[ndim-1].idx != NULL) { + index = q[ndim-1].idx; + } else { + //index = ALLOC_N(size_t, size); + index = (size_t*)cumo_cuda_runtime_malloc(sizeof(size_t)*size); + use_cumo_cuda_runtime_malloc = 1; + } + CUMO_SDX_SET_INDEX(na2->stridx[0], index); + + for (i=ndim-1; i>=0; i--) { + cumo_stridx_t sdx1 = na1->stridx[q[i].orig_dim]; + if (i==ndim-1) { + if (size == 0) { + rb_raise(cumo_na_eShapeError, "cannot get element of empty array"); + } + } else { + if (size != q[i].n) { + rb_raise(cumo_na_eShapeError, "index array sizes mismatch"); + } + } + + if (q[i].idx != NULL && CUMO_SDX_IS_INDEX(sdx1)) { + // index <- index + size_t *index1 = CUMO_SDX_GET_INDEX(sdx1); + if (i==ndim-1) { + cumo_na_index_aref_naview_index_index_kernel_launch(index, index1, size); + } else { + cumo_na_index_at_naview_index_index_index_add_kernel_launch(index, index1, q[i].idx, size); + } + q[i].idx = NULL; + } + else if (q[i].idx == NULL && CUMO_SDX_IS_INDEX(sdx1)) { + // step <- index + size_t beg = q[i].beg; + ssize_t step = q[i].step; + size_t *index1 = CUMO_SDX_GET_INDEX(sdx1); + if (i==ndim-1) { + cumo_na_index_aref_naview_index_index_beg_step_kernel_launch(index, index1, beg, step, size); + } else { + cumo_na_index_at_naview_index_index_beg_step_add_kernel_launch(index, index1, beg, step, size); + } + } + else if (q[i].idx != NULL && CUMO_SDX_IS_STRIDE(sdx1)) { + // index <- step + ssize_t stride1 = CUMO_SDX_GET_STRIDE(sdx1); + if (stride1<0) { + size_t last; + stride1 = -stride1; + last = na1->base.shape[q[i].orig_dim] - 1; + if (na2->offset < last * stride1) { + rb_raise(rb_eStandardError,"bug: negative offset"); + } + na2->offset -= last * stride1; + if (i==ndim-1) { + cumo_na_index_aref_naview_index_stride_last_kernel_launch(index, stride1, last, size); + } else { + cumo_na_index_at_naview_index_stride_last_add_kernel_launch(index, stride1, last, size); + } + } else { + if (i==ndim-1) { + cumo_na_index_aref_nadata_index_stride_kernel_launch(index, stride1, size); + } else { + cumo_na_index_at_nadata_index_stride_add_kernel_launch(index, q[i].idx, stride1, size); + } + } + q[i].idx = NULL; + } + else if (q[i].idx == NULL && CUMO_SDX_IS_STRIDE(sdx1)) { + // step <- step + size_t beg = q[i].beg; + ssize_t step = q[i].step; + ssize_t stride1 = CUMO_SDX_GET_STRIDE(sdx1); + if (stride1<0) { + size_t last; + stride1 = -stride1; + last = na1->base.shape[q[i].orig_dim] - 1; + if (na2->offset < last * stride1) { + rb_raise(rb_eStandardError,"bug: negative offset"); + } + na2->offset -= last * stride1; + if (i==ndim-1) { + cumo_na_index_at_nadata_index_beg_step_stride_kernel_launch(index, last - beg, -step, stride1, size); + } else { + cumo_na_index_at_nadata_index_beg_step_stride_add_kernel_launch(index, last - beg, -step, stride1, size); + } + } else { + if (i==ndim-1) { + cumo_na_index_at_nadata_index_beg_step_stride_kernel_launch(index, beg, step, stride1, size); + } else { + cumo_na_index_at_nadata_index_beg_step_stride_add_kernel_launch(index, beg, step, stride1, size); + } + } + } + } + na2->base.size = size; + na2->base.shape[0] = size; + if (use_cumo_cuda_runtime_malloc) { + CUMO_SHOW_SYNCHRONIZE_FIXME_WARNING_ONCE("index", "cumo_na_index_at_naview"); + cumo_cuda_runtime_check_status(cudaDeviceSynchronize()); + } +} + static int cumo_na_ndim_new_narray(int ndim, const cumo_na_index_arg_t *q) { @@ -587,6 +768,7 @@ typedef struct { cumo_narray_t *na1; int keep_dim; size_t pos; // offset position for 0-dimensional narray. 0-dimensional array does not use q. + int at_mode; // 0: aref, 1: at } cumo_na_aref_md_data_t; static cumo_na_index_arg_t* @@ -614,6 +796,7 @@ VALUE cumo_na_aref_md_protected(VALUE data_value) cumo_na_index_arg_t *q = data->q; cumo_narray_t *na1 = data->na1; int keep_dim = data->keep_dim; + int at_mode = data->at_mode; int ndim_new; VALUE view; @@ -624,10 +807,14 @@ VALUE cumo_na_aref_md_protected(VALUE data_value) if (cumo_na_debug_flag) print_index_arg(q,ndim); - if (keep_dim) { - ndim_new = ndim; + if (at_mode) { + ndim_new = 1; } else { - ndim_new = cumo_na_ndim_new_narray(ndim, q); + if (keep_dim) { + ndim_new = ndim; + } else { + ndim_new = cumo_na_ndim_new_narray(ndim, q); + } } view = cumo_na_s_allocate_view(rb_obj_class(self)); @@ -647,7 +834,11 @@ VALUE cumo_na_aref_md_protected(VALUE data_value) na2->offset = data->pos; na2->base.size = 1; } else { - cumo_na_index_aref_nadata((cumo_narray_data_t *)na1,na2,q,elmsz,ndim,keep_dim); + if (at_mode) { + cumo_na_index_at_nadata((cumo_narray_data_t *)na1,na2,q,elmsz,ndim,keep_dim); + } else { + cumo_na_index_aref_nadata((cumo_narray_data_t *)na1,na2,q,elmsz,ndim,keep_dim); + } } na2->data = self; break; @@ -659,7 +850,11 @@ VALUE cumo_na_aref_md_protected(VALUE data_value) } else { na2->offset = ((cumo_narray_view_t *)na1)->offset; na2->data = ((cumo_narray_view_t *)na1)->data; - cumo_na_index_aref_naview((cumo_narray_view_t *)na1,na2,q,elmsz,ndim,keep_dim); + if (at_mode) { + cumo_na_index_at_naview((cumo_narray_view_t *)na1,na2,q,elmsz,ndim,keep_dim); + } else { + cumo_na_index_aref_naview((cumo_narray_view_t *)na1,na2,q,elmsz,ndim,keep_dim); + } } break; } @@ -684,7 +879,7 @@ cumo_na_aref_md_ensure(VALUE data_value) } static VALUE -cumo_na_aref_md(int argc, VALUE *argv, VALUE self, int keep_dim, int result_nd, size_t pos) +cumo_na_aref_md(int argc, VALUE *argv, VALUE self, int keep_dim, int result_nd, size_t pos, int at_mode) { VALUE args; // should be GC protected cumo_narray_t *na1; @@ -696,6 +891,9 @@ cumo_na_aref_md(int argc, VALUE *argv, VALUE self, int keep_dim, int result_nd, CumoGetNArray(self,na1); args = rb_ary_new4(argc,argv); + if (at_mode && na1->ndim == 0) { + rb_raise(cumo_na_eDimensionError,"argument length does not match dimension size"); + } if (argc == 1 && result_nd == 1) { idx = argv[0]; @@ -724,6 +922,7 @@ cumo_na_aref_md(int argc, VALUE *argv, VALUE self, int keep_dim, int result_nd, data.q = cumo_na_allocate_index_args(result_nd); data.na1 = na1; data.keep_dim = keep_dim; + data.at_mode = at_mode; switch(na1->type) { case CUMO_NARRAY_DATA_T: @@ -760,7 +959,15 @@ cumo_na_aref_main(int nidx, VALUE *idx, VALUE self, int keep_dim, int result_nd, return rb_funcall(*idx,cumo_id_mask,1,self); } } - return cumo_na_aref_md(nidx, idx, self, keep_dim, result_nd, pos); + return cumo_na_aref_md(nidx, idx, self, keep_dim, result_nd, pos, 0); +} + +/* method: at([idx1,idx2,...,idxN], [idx1,idx2,...,idxN]) */ +VALUE +cumo_na_at_main(int nidx, VALUE *idx, VALUE self, int keep_dim, int result_nd, size_t pos) +{ + cumo_na_index_arg_to_internal_order(nidx, idx, self); + return cumo_na_aref_md(nidx, idx, self, keep_dim, result_nd, pos, 1); } diff --git a/ext/cumo/narray/index_kernel.cu b/ext/cumo/narray/index_kernel.cu index 6a6357d1..6cdd294f 100644 --- a/ext/cumo/narray/index_kernel.cu +++ b/ext/cumo/narray/index_kernel.cu @@ -42,6 +42,48 @@ __global__ void cumo_na_index_aref_naview_index_index_beg_step_kernel(size_t *id } } +__global__ void cumo_na_index_at_nadata_index_beg_step_stride_kernel(size_t *idx, size_t beg, ssize_t step, ssize_t s1, uint64_t n) +{ + for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) { + idx[i] = (beg + step * i) * s1; + } +} + +__global__ void cumo_na_index_at_nadata_index_beg_step_stride_add_kernel(size_t *idx, size_t beg, ssize_t step, ssize_t s1, uint64_t n) +{ + for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) { + idx[i] += (beg + step * i) * s1; + } +} + +__global__ void cumo_na_index_at_nadata_index_stride_add_kernel(size_t *idx, size_t *idx1, ssize_t s1, uint64_t n) +{ + for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) { + idx[i] += idx1[i] * s1; + } +} + +__global__ void cumo_na_index_at_naview_index_index_index_add_kernel(size_t *idx, size_t *idx1, size_t *idx2, uint64_t n) +{ + for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) { + idx[i] += idx1[idx2[i]]; + } +} + +__global__ void cumo_na_index_at_naview_index_index_beg_step_add_kernel(size_t *idx, size_t *idx1, size_t beg, ssize_t step, uint64_t n) +{ + for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) { + idx[i] += idx1[beg + step * i]; + } +} + +__global__ void cumo_na_index_at_naview_index_stride_last_add_kernel(size_t *idx, ssize_t s1, size_t last, uint64_t n) +{ + for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) { + idx[i] += (last - idx[i]) * s1; + } +} + void cumo_na_index_aref_nadata_index_stride_kernel_launch(size_t *idx, ssize_t s1, uint64_t n) { size_t grid_dim = cumo_get_grid_dim(n); @@ -77,6 +119,48 @@ void cumo_na_index_aref_naview_index_index_beg_step_kernel_launch(size_t *idx, s cumo_na_index_aref_naview_index_index_beg_step_kernel<<>>(idx, idx1, beg, step, n); } +void cumo_na_index_at_nadata_index_stride_add_kernel_launch(size_t *idx, size_t *idx1, ssize_t s1, uint64_t n) +{ + size_t grid_dim = cumo_get_grid_dim(n); + size_t block_dim = cumo_get_block_dim(n); + cumo_na_index_at_nadata_index_stride_add_kernel<<>>(idx, idx1, s1, n); +} + +void cumo_na_index_at_nadata_index_beg_step_stride_kernel_launch(size_t *idx, size_t beg, ssize_t step, ssize_t s1, uint64_t n) +{ + size_t grid_dim = cumo_get_grid_dim(n); + size_t block_dim = cumo_get_block_dim(n); + cumo_na_index_at_nadata_index_beg_step_stride_kernel<<>>(idx, beg, step, s1, n); +} + +void cumo_na_index_at_nadata_index_beg_step_stride_add_kernel_launch(size_t *idx, size_t beg, ssize_t step, ssize_t s1, uint64_t n) +{ + size_t grid_dim = cumo_get_grid_dim(n); + size_t block_dim = cumo_get_block_dim(n); + cumo_na_index_at_nadata_index_beg_step_stride_add_kernel<<>>(idx, beg, step, s1, n); +} + +void cumo_na_index_at_naview_index_index_index_add_kernel_launch(size_t *idx, size_t *idx1, size_t *idx2, uint64_t n) +{ + size_t grid_dim = cumo_get_grid_dim(n); + size_t block_dim = cumo_get_block_dim(n); + cumo_na_index_at_naview_index_index_index_add_kernel<<>>(idx, idx1, idx2, n); +} + +void cumo_na_index_at_naview_index_index_beg_step_add_kernel_launch(size_t *idx, size_t *idx1, size_t beg, ssize_t step, uint64_t n) +{ + size_t grid_dim = cumo_get_grid_dim(n); + size_t block_dim = cumo_get_block_dim(n); + cumo_na_index_at_naview_index_index_beg_step_add_kernel<<>>(idx, idx1, beg, step, n); +} + +void cumo_na_index_at_naview_index_stride_last_add_kernel_launch(size_t *idx, ssize_t s1, size_t last, uint64_t n) +{ + size_t grid_dim = cumo_get_grid_dim(n); + size_t block_dim = cumo_get_block_dim(n); + cumo_na_index_at_naview_index_stride_last_add_kernel<<>>(idx, s1, last, n); +} + #if defined(__cplusplus) #if 0 { /* satisfy cc-mode */ diff --git a/lib/cumo/narray/extra.rb b/lib/cumo/narray/extra.rb index f22ba06b..6e789990 100644 --- a/lib/cumo/narray/extra.rb +++ b/lib/cumo/narray/extra.rb @@ -43,56 +43,6 @@ def flipud reverse(0) end - # Multi-dimensional array indexing. - # Same as [] for one-dimensional NArray. - # Similar to numpy's tuple indexing, i.e., `a[[1,2,..],[3,4,..]]` - # (This method will be rewritten in C) - # @return [Cumo::NArray] one-dimensional view of self. - # @example - # p x = Cumo::DFloat.new(3,3,3).seq - # # Cumo::DFloat#shape=[3,3,3] - # # [[[0, 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]]] - # - # p x.at([0,1,2],[0,1,2],[-1,-2,-3]) - # # Cumo::DFloat(view)#shape=[3] - # # [2, 13, 24] - def at(*indices) - if indices.size != ndim - raise DimensionError, "argument length does not match dimension size" - end - idx = nil - stride = 1 - (indices.size-1).downto(0) do |i| - ix = Int64.cast(indices[i]) - if ix.ndim != 1 - raise DimensionError, "index array is not one-dimensional" - end - ix[ix < 0] += shape[i] - if ((ix < 0) & (ix >= shape[i])).any? - raise IndexError, "index array is out of range" - end - if idx - if idx.size != ix.size - raise ShapeError, "index array sizes mismatch" - end - idx += ix * stride - stride *= shape[i] - else - idx = ix - stride = shape[i] - end - end - self[idx] - end - # Rotate in the plane specified by axes. # @example # p a = Cumo::Int32.new(2,2).seq diff --git a/test/narray_test.rb b/test/narray_test.rb index 433a46d5..5152e831 100644 --- a/test/narray_test.rb +++ b/test/narray_test.rb @@ -73,6 +73,19 @@ def setup assert { a[5] == 11 } assert { a[5].size == 1 } assert { a[-1] == 11 } + + assert { a.at([3, 4]) == [5,7] } + assert { a.view.at([3, 4]) == [5,7] } + assert { a[2..-1].at([1, 2]) == [5,7] } + assert { a.at(Cumo::Int32.cast([3, 4])) == [5,7] } + assert { a.view.at(Cumo::Int32.cast([3, 4])) == [5,7] } + assert { a.at(3..4) == [5,7] } + assert { a.view.at(3..4) == [5,7] } + assert { a.at([5]) == [11] } + assert { a.view.at([5]) == [11] } + assert { a.at([-1]) == [11] } + assert { a.view.at([-1]) == [11] } + assert { a[(0..-1).each] == [1,2,3,5,7,11] } assert { a[(0...-1).each] == [1,2,3,5,7] } @@ -226,6 +239,20 @@ def setup assert { a[1,2] == src[1][2] } assert { a[3..4] == [5,7] } assert { a[0,1..2] == [2,3] } + + assert { a.at([0,1],[1,2]) == [2,11] } + assert { a.view.at([0,1],[1,2]) == [2,11] } + assert { a.at([0,1],(0..2) % 2) == [1,11] } + assert { a.view.at([0,1],(0..2) % 2) == [1,11] } + assert { a.at((0..1) % 1,[0,2]) == [1,11] } + assert { a.view.at((0..1) % 1,[0,2]) == [1,11] } + assert { a.at(Cumo::Int32.cast([0,1]),Cumo::Int32.cast([1,2])) == [2,11] } + assert { a.view.at(Cumo::Int32.cast([0,1]),Cumo::Int32.cast([1,2])) == [2,11] } + assert { a[[0,1],[0,2]].at([0,1],[0,1]) == [1,11] } + assert { a[[0,1],(0..2) % 2].at([0,1],[0,1]) == [1,11] } + assert { a[(0..1) % 1,[0,2]].at([0,1],[0,1]) == [1,11] } + assert { a[(0..1) % 1,(0..2) % 2].at([0,1],[0,1]) == [1,11] } + assert { a[0,:*] == src[0] } assert { a[1,:*] == src[1] } assert { a[:*,1] == [src[0][1],src[1][1]] } @@ -400,6 +427,9 @@ def setup assert_raise(IndexError) { a[1, 1, 1, :rest, 1] } assert_raise(IndexError) { a[:rest, 1, :rest, 0] } + assert { a.at([0,1],[1,0], [0,1]) == [3,6] } + assert { a.view.at([0,1],[1,0],[0,1]) == [3,6] } + assert { a.transpose == [[[1,5],[3,7]],[[2,6],[4,8]]] } assert { a.transpose(2,1,0) == [[[1,5],[3,7]],[[2,6],[4,8]]] } assert { a.transpose(0,2,1) == [[[1,3],[2,4]],[[5,7],[6,8]]] } @@ -637,6 +667,11 @@ def setup diag = a.dup[[0,1],[0,1]].diagonal diag.inplace - 1 assert { diag == [0, 4] } + + assert { a.at([0,1],[0,1]).dup == [1, 5] } + at = a.dup + at.at([0,1],[0,1]).inplace - 1 + assert { at == [[0,2,3],[4,4,6]] } end end end