<commit_before>import unittest

import cupy
from cupy.random import distributions
from cupy import testing


@testing.parameterize(*testing.product({
    'shape': [(4, 3, 2), (3, 2)],
    'n_shape': [(), (3, 2)],
    'p_shape': [(), (3, 2)],
})
)
@testing.gpu
class TestDistributionsBinomial(unittest.TestCase):

    def check_distribution(self, dist_func, n_dtype, p_dtype, dtype):
        n = 5 * cupy.ones(self.n_shape, dtype=n_dtype)
        p = 0.5 * cupy.ones(self.p_shape, dtype=p_dtype)
        out = dist_func(n, p, self.shape, dtype)
        self.assertEqual(self.shape, out.shape)
        self.assertEqual(out.dtype, dtype)

    @cupy.testing.for_int_dtypes('dtype')
    @cupy.testing.for_int_dtypes('n_dtype')
    @cupy.testing.for_float_dtypes('p_dtype')
    def test_binomial(self, n_dtype, p_dtype, dtype):
        self.check_distribution(distributions.binomial,
                                n_dtype, p_dtype, dtype)


@testing.parameterize(*testing.product({
    'shape': [(4, 3, 2), (3, 2)],
    'loc_shape': [(), (3, 2)],
    'scale_shape': [(), (3, 2)],
})
)
@testing.gpu
class TestDistributionsGumbel(unittest.TestCase):

    def check_distribution(self, dist_func, loc_dtype, scale_dtype, dtype):
        loc = cupy.ones(self.loc_shape, dtype=loc_dtype)
        scale = cupy.ones(self.scale_shape, dtype=scale_dtype)
        out = dist_func(loc, scale, self.shape, dtype)
        self.assertEqual(self.shape, out.shape)
        self.assertEqual(out.dtype, dtype)

    @cupy.testing.for_float_dtypes('dtype', no_float16=True)
    @cupy.testing.for_float_dtypes('loc_dtype')
    @cupy.testing.for_float_dtypes('scale_dtype')
    def test_gumbel(self, loc_dtype, scale_dtype, dtype):
        self.check_distribution(distributions.gumbel,
                                loc_dtype, scale_dtype, dtype)

<commit_msg>[fix] escape timeout<commit_after>import unittest

import cupy
from cupy.random import distributions
from cupy import testing

import numpy


_signed_dtypes = tuple(numpy.dtype(i).type for i in 'bhilq')
_unsigned_dtypes = tuple(numpy.dtype(i).type for i in 'BHILQ')
_int_dtypes = _signed_dtypes + _unsigned_dtypes


@testing.parameterize(*testing.product({
    'shape': [(4, 3, 2), (3, 2)],
    'n_shape': [(), (3, 2)],
    'p_shape': [(), (3, 2)],
    'dtype': _int_dtypes,  # to escape timeout
})
)
@testing.gpu
class TestDistributionsBinomial(unittest.TestCase):

    def check_distribution(self, dist_func, n_dtype, p_dtype, dtype):
        n = 5 * cupy.ones(self.n_shape, dtype=n_dtype)
        p = 0.5 * cupy.ones(self.p_shape, dtype=p_dtype)
        out = dist_func(n, p, self.shape, dtype)
        self.assertEqual(self.shape, out.shape)
        self.assertEqual(out.dtype, dtype)

    @cupy.testing.for_int_dtypes('n_dtype')
    @cupy.testing.for_float_dtypes('p_dtype')
    def test_binomial(self, n_dtype, p_dtype):
        self.check_distribution(distributions.binomial,
                                n_dtype, p_dtype, self.dtype)


@testing.parameterize(*testing.product({
    'shape': [(4, 3, 2), (3, 2)],
    'loc_shape': [(), (3, 2)],
    'scale_shape': [(), (3, 2)],
})
)
@testing.gpu
class TestDistributionsGumbel(unittest.TestCase):

    def check_distribution(self, dist_func, loc_dtype, scale_dtype, dtype):
        loc = cupy.ones(self.loc_shape, dtype=loc_dtype)
        scale = cupy.ones(self.scale_shape, dtype=scale_dtype)
        out = dist_func(loc, scale, self.shape, dtype)
        self.assertEqual(self.shape, out.shape)
        self.assertEqual(out.dtype, dtype)

    @cupy.testing.for_float_dtypes('dtype', no_float16=True)
    @cupy.testing.for_float_dtypes('loc_dtype')
    @cupy.testing.for_float_dtypes('scale_dtype')
    def test_gumbel(self, loc_dtype, scale_dtype, dtype):
        self.check_distribution(distributions.gumbel,
                                loc_dtype, scale_dtype, dtype)


{'fraction_non_alphanumeric': 0.14115822130299896, 'fraction_numerical': 0.011892450879007239, 'mean_word_length': 3.5357561547479484, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4148092', 'n_tokens_mistral': 1407, 'n_tokens_neox': 1384, 'n_words': 248}