diff --git a/doc/array.rst b/doc/array.rst index 34c65693e8c4244c5a6c401eb35d4573a56b0971..dacb6cbf46d88f6b6b7d07729df6afc4c4be42c7 100644 --- a/doc/array.rst +++ b/doc/array.rst @@ -293,6 +293,11 @@ Constructing :class:`GPUArray` Instances Same as :func:`empty`, but the :class:`GPUArray` is zero-initialized before being returned. +.. function:: ones(shape, dtype=np.float64, *, allocator=None, order="C") + + Same as :func:`empty`, but the :class:`GPUArray` is one-initialized before + being returned. + .. function:: empty_like(other_ary, dtype=None, order="K") Make a new, uninitialized :class:`GPUArray` having the same properties diff --git a/pycuda/gpuarray.py b/pycuda/gpuarray.py index a3dbc100738396b897c5699c0de852e1bee92aa1..e6238231c26f2b2243c3fcfeebbd5318b07ee43f 100644 --- a/pycuda/gpuarray.py +++ b/pycuda/gpuarray.py @@ -1248,6 +1248,14 @@ def zeros(shape, dtype, allocator=drv.mem_alloc, order="C"): return result +def ones(shape, dtype=np.float64, allocator=drv.mem_alloc, order="C"): + """Returns an array of the given shape and dtype filled with 1's.""" + result = GPUArray(shape, dtype, allocator, order=order) + one = np.ones((), dtype) + result.fill(one) + return result + + def _array_like_helper(other_ary, dtype, order): """Set order, strides, dtype as in numpy's zero_like. """ strides = None diff --git a/test/test_gpuarray.py b/test/test_gpuarray.py index 73ec3ade3f99e8830882b6271aadd70945e4b4f3..ddd8b2f6eb0fac4a47795336eec35affdd4a0699 100644 --- a/test/test_gpuarray.py +++ b/test/test_gpuarray.py @@ -485,6 +485,15 @@ class TestGPUArray: a = gpuarray.arange(12, dtype=np.float32) assert (np.arange(12, dtype=np.float32) == a.get()).all() + @mark_cuda_test + def test_ones(self): + + ones = np.ones(10) + ones_gpu = gpuarray.ones(10) + + np.testing.assert_allclose(ones, ones_gpu.get(), rtol=1e-6) + assert ones.dtype == ones_gpu.dtype + @mark_cuda_test def test_stack(self):