lebedov / scikit-cuda

Python interface to GPU-powered libraries
http://scikit-cuda.readthedocs.org/
Other
973 stars 179 forks source link

skcuda.misc.sum for >2 dimensional arrays #190

Open biochem-fan opened 7 years ago

biochem-fan commented 7 years ago

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.

stevendbrown commented 6 years ago

I get similar unexpected behavior in skcuda.misc.var.

stevendbrown commented 6 years ago

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.