Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 26 additions & 27 deletions mkl_random/mklrand.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -1794,6 +1794,22 @@ def _brng_id_to_name(int brng_id):
return brng_name


# Largest half-open interval [low, high) per randint dtype.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

missing a changelog entry

_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
Expand Down Expand Up @@ -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]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

key now holds the canonical dtype name, not the dict key, so name would be clearer

except KeyError:
raise TypeError(
f'Unsupported dtype "{_dtype.name}" for randint'
) from None

return (lowbnd, highbnd,
getattr(self, f"_rand_{key}"),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As a possible possible follow-up: each call still builds two method names with f-strings and looks them up with getattr, so storing the method names in the table would remove that.

getattr(self, f"_rand_{key}_broadcast"))
Expand Down
Loading