Merge pull request #8115 from kmaehashi/fix-generator-api-slow

Fix slow test
cupy · Jan 17, 2024 · 09395a1 · 09395a1
2 parents 9016b92 + 74da76f
commit 09395a1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/cupy/random/__init__.py b/cupy/random/__init__.py
@@ -50,9 +50,10 @@ def __getattr__(key):
     # users use the class for type annotation.
     if key == 'Generator':
         # Lazy import libraries depending on cuRAND
-        import cupy.random._generator_api as _generator_api
-        _cupy.random.Generator = _generator_api.Generator
-        return _generator_api.Generator
+        import cupy.random._generator_api
+        Generator = cupy.random._generator_api.Generator
+        _cupy.random.Generator = Generator
+        return Generator
     raise AttributeError(f"module '{__name__}' has no attribute '{key}'")
 
 

diff --git a/tests/cupy_tests/random_tests/test_generator_api.py b/tests/cupy_tests/random_tests/test_generator_api.py
@@ -22,7 +22,7 @@ class GeneratorTestCase(common_distributions.BaseGeneratorTestCase):
 
     def get_rng(self, xp, seed):
         if xp is cupy:
-            return cupy.random._generator_api.Generator(
+            return cupy.random.Generator(
                 random._bit_generator.Philox4x3210(seed=seed))
         else:
             return numpy.random.Generator(numpy.random.MT19937(seed))