scikit-cuda
scikit-cuda copied to clipboard
skcuda.misc.sum for >2 dimensional arrays
Hi, it seems that skcuda.misc.sum does not support arrays with more than two dimensions.
For example, the following code does not work as expected.
# based on https://github.com/lebedov/scikit-cuda/issues/173
import skcuda.misc
skcuda.misc.init()
dev = skcuda.misc.init_device()
skcuda.misc.init_context(dev)
a = skcuda.misc.ones((2, 3, 4), 'float64')
skcuda.misc.sum(a, axis=0) # should return a (3, 4) array with all elements = 2
# but returns [2, 2, 2]
skcuda.misc.sum(a, axis=1) # should return a (2, 4) array with all elements = 3
# but returns [3, 3]
skcuda.misc.sum(a, axis=2) # should return a (2, 3) array with all elements = 4
# but says "invalid axis"
For the special case of axis=0, a workaround is to reshape the array before summing.
I would appreciate a lot if this operation could be supported by the library.
I get similar unexpected behavior in skcuda.misc.var.
An axis=0 reshape workaround is also effective for skcuda.misc.var, which conveniently allows for easy reduction of the GPU memory burden required for var by splitting the reshaped array into chunks on-GPU, computing var, downloading to the host, and then reshaping in host RAM once all var calls are complete.