diff --git a/mkl_random/mklrand.pyx b/mkl_random/mklrand.pyx index f8d09c1..e302da3 100644 --- a/mkl_random/mklrand.pyx +++ b/mkl_random/mklrand.pyx @@ -1794,6 +1794,22 @@ def _brng_id_to_name(int brng_id): return brng_name +# Largest half-open interval [low, high) per randint dtype. +_randint_bounds = { + np.dtype(name): (name, low, high) for name, low, high in ( + ("bool", 0, 2), + ("int8", -2**7, 2**7), + ("int16", -2**15, 2**15), + ("int32", -2**31, 2**31), + ("int64", -2**63, 2**63), + ("uint8", 0, 2**8), + ("uint16", 0, 2**16), + ("uint32", 0, 2**32), + ("uint64", 0, 2**64), + ) +} + + cdef class _MKLRandomState: cdef irk_state *internal_state cdef object lock @@ -2143,35 +2159,18 @@ cdef class _MKLRandomState: self.internal_state, irk_double_vec, size, self.lock ) - # Set up dictionary of integer types and relevant functions. - # - # The dictionary is keyed by dtype(...).name and the values - # are a tuple (low, high, function), where low and high are - # the bounds of the largest half open interval `[low, high)` - # and the function is the relevant function to call for - # that precision. - # - # The functions are all the same except for changed types in - # a few places. It would be easy to template them. + # Returns randint bounds and the matching `_rand_*`/`_rand_*_broadcast` + # methods for dtype. def _choose_randint_type(self, dtype): - _randint_bounds = { - "bool": (0, 2), - "int8": (-2**7, 2**7), - "int16": (-2**15, 2**15), - "int32": (-2**31, 2**31), - "int64": (-2**63, 2**63), - "uint8": (0, 2**8), - "uint16": (0, 2**16), - "uint32": (0, 2**32), - "uint64": (0, 2**64), - } - - key = np.dtype(dtype).name - if key not in _randint_bounds: - raise TypeError(f'Unsupported dtype "{key}" for randint') - - lowbnd, highbnd = _randint_bounds[key] + _dtype = np.dtype(dtype) + try: + key, lowbnd, highbnd = _randint_bounds[_dtype] + except KeyError: + raise TypeError( + f'Unsupported dtype "{_dtype.name}" for randint' + ) from None + return (lowbnd, highbnd, getattr(self, f"_rand_{key}"), getattr(self, f"_rand_{key}_broadcast"))