@@ -337,9 +337,107 @@ def test_bincount_too_small_minlength(self, dtype):
337337 xp .bincount (x , minlength = - 1 )
338338
339339
340- # TODO(leofang): we temporarily remove CUB histogram support for now,
341- # see cupy/cupy#7698. When it's ready, revert the commit that checked
342- # in this comment to restore the support.
340+ # This class compares CUB results against NumPy's
341+ @unittest .skipUnless (False , "The CUB routine is not enabled" )
342+ class TestCubHistogram (unittest .TestCase ):
343+
344+ def setUp (self ):
345+ self .old_accelerators = _accelerator .get_routine_accelerators ()
346+ _accelerator .set_routine_accelerators (["cub" ])
347+
348+ def tearDown (self ):
349+ _accelerator .set_routine_accelerators (self .old_accelerators )
350+
351+ @testing .for_all_dtypes (no_bool = True , no_complex = True )
352+ @testing .numpy_cupy_array_equal ()
353+ def test_histogram (self , xp , dtype ):
354+ x = testing .shaped_arange ((10 ,), xp , dtype )
355+
356+ if xp is numpy :
357+ return xp .histogram (x )
358+
359+ # xp is cupy, first ensure we really use CUB
360+ cub_func = "cupy._statistics.histogram.cub.device_histogram"
361+ with testing .AssertFunctionIsCalled (cub_func ):
362+ xp .histogram (x )
363+ # ...then perform the actual computation
364+ return xp .histogram (x )
365+
366+ @testing .for_all_dtypes (no_bool = True , no_complex = True )
367+ @testing .numpy_cupy_array_equal ()
368+ def test_histogram_range_float (self , xp , dtype ):
369+ a = testing .shaped_arange ((10 ,), xp , dtype )
370+ h , b = xp .histogram (a , testing .shaped_arange ((10 ,), xp , numpy .float64 ))
371+ assert int (h .sum ()) == 10
372+ return h , b
373+
374+ @testing .for_all_dtypes_combination (
375+ ["dtype_a" , "dtype_b" ], no_bool = True , no_complex = True
376+ )
377+ @testing .numpy_cupy_array_equal ()
378+ def test_histogram_with_bins (self , xp , dtype_a , dtype_b ):
379+ x = testing .shaped_arange ((10 ,), xp , dtype_a )
380+ bins = testing .shaped_arange ((4 ,), xp , dtype_b )
381+
382+ if xp is numpy :
383+ return xp .histogram (x , bins )[0 ]
384+
385+ # xp is cupy, first ensure we really use CUB
386+ cub_func = "cupy._statistics.histogram.cub.device_histogram"
387+ with testing .AssertFunctionIsCalled (cub_func ):
388+ xp .histogram (x , bins )
389+ # ...then perform the actual computation
390+ return xp .histogram (x , bins )[0 ]
391+
392+ @testing .for_all_dtypes_combination (
393+ ["dtype_a" , "dtype_b" ], no_bool = True , no_complex = True
394+ )
395+ @testing .numpy_cupy_array_equal ()
396+ def test_histogram_with_bins2 (self , xp , dtype_a , dtype_b ):
397+ x = testing .shaped_arange ((10 ,), xp , dtype_a )
398+ bins = testing .shaped_arange ((4 ,), xp , dtype_b )
399+
400+ if xp is numpy :
401+ return xp .histogram (x , bins )[1 ]
402+
403+ # xp is cupy, first ensure we really use CUB
404+ cub_func = "cupy._statistics.histogram.cub.device_histogram"
405+ with testing .AssertFunctionIsCalled (cub_func ):
406+ xp .histogram (x , bins )
407+ # ...then perform the actual computation
408+ return xp .histogram (x , bins )[1 ]
409+
410+ @testing .slow
411+ @testing .numpy_cupy_array_equal ()
412+ def test_no_oom (self , xp ):
413+ # ensure the workaround for NVIDIA/cub#613 kicks in
414+ amax = 28854312
415+ A = xp .linspace (
416+ 0 , amax , num = amax , endpoint = True , retstep = False , dtype = xp .int32
417+ )
418+ out = xp .histogram (A , bins = amax , range = [0 , amax ])
419+ return out
420+
421+ @testing .for_int_dtypes ("dtype" , no_bool = True )
422+ @testing .numpy_cupy_array_equal ()
423+ def test_bincount_gh7698 (self , xp , dtype ):
424+ dtype = xp .dtype (dtype )
425+ max_val = xp .iinfo (dtype ).max if dtype .itemsize < 4 else 65536
426+ if dtype == xp .uint64 :
427+ pytest .skip ("only numpy raises exception on uint64 input" )
428+
429+ # https://github.com/cupy/cupy/issues/7698
430+ x = xp .arange (max_val , dtype = dtype )
431+
432+ if xp is numpy :
433+ return xp .bincount (x )
434+
435+ # xp is cupy, first ensure we really use CUB
436+ cub_func = "cupy._statistics.histogram.cub.device_histogram"
437+ with testing .AssertFunctionIsCalled (cub_func ):
438+ xp .bincount (x )
439+ # ...then perform the actual computation
440+ return xp .bincount (x )
343441
344442
345443@testing .parameterize (
@@ -560,7 +658,7 @@ def test_histogramdd_disallow_arraylike_bins(self):
560658 * testing .product (
561659 {
562660 "weights" : [None , 1 , 2 ],
563- "weights_dtype" : [numpy .int32 , numpy . float32 ],
661+ "weights_dtype" : [numpy .int32 , cupy . default_float_type () ],
564662 "density" : [True , False ],
565663 "bins" : [10 , (8 , 16 ), (16 , 8 ), "array_list" , "array" ],
566664 "range" : [None , ((20 , 50 ), (10 , 100 ))],
0 commit comments