From d9ddc239315439a630d36788d86d8bf9924c5323 Mon Sep 17 00:00:00 2001 From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com> Date: Mon, 1 Jun 2026 09:46:03 +0200 Subject: [PATCH 1/2] Validate PyTorch random size arguments --- src/pyrecest/_backend/pytorch/random.py | 62 ++++++++++++++----------- 1 file changed, 35 insertions(+), 27 deletions(-) diff --git a/src/pyrecest/_backend/pytorch/random.py b/src/pyrecest/_backend/pytorch/random.py index 6d9b0bfd5..0d6a7d4d0 100644 --- a/src/pyrecest/_backend/pytorch/random.py +++ b/src/pyrecest/_backend/pytorch/random.py @@ -1,5 +1,6 @@ """Torch based random backend.""" +from math import prod as _prod from numbers import Integral as _Integral import torch as _torch @@ -17,21 +18,42 @@ } +def _size_type_error(): + return TypeError("size must be None, an integer, or a sequence of integers") + + +def _looks_like_integer_dimension(value): + return isinstance(value, _Integral) and not isinstance(value, bool) + + +def _integer_dimension(value): + if not _looks_like_integer_dimension(value): + raise _size_type_error() + value = int(value) + if value < 0: + raise ValueError("size dimensions must be non-negative") + return value + + +def _shape_from_size(size): + if size is None: + return () + if _looks_like_integer_dimension(size): + return (_integer_dimension(size),) + if isinstance(size, (str, bytes)) or not hasattr(size, "__iter__"): + raise _size_type_error() + return tuple(_integer_dimension(dim) for dim in size) + + def _choice_size(size): if size is None: return None, 1 - if not hasattr(size, "__iter__"): - size = (size,) - size = tuple(int(dim) for dim in size) - return size, int(_torch.prod(_torch.tensor(size)).item()) + size = _shape_from_size(size) + return size, _prod(size) if size else 1 def _randint_size(size): - if size is None: - return () - if not hasattr(size, "__iter__") or isinstance(size, (str, bytes)): - return (size,) - return tuple(size) + return _shape_from_size(size) def randint(low, high=None, size=None, *args, **kwargs): @@ -45,9 +67,7 @@ def randint(low, high=None, size=None, *args, **kwargs): def _normal_size(size): if size is None: return None - if not hasattr(size, "__iter__"): - return (size,) - return tuple(int(dim) for dim in size) + return _shape_from_size(size) def _normal_device(*values): @@ -166,11 +186,7 @@ def seed(*args, **kwargs): def rand(size=None, dtype=None): - if size is None: - size = () - elif not hasattr(size, "__iter__"): - size = (size,) - return _torch.rand(size, dtype=dtype) + return _torch.rand(_shape_from_size(size), dtype=dtype) def multinomial(n, pvals): @@ -207,11 +223,7 @@ def uniform(low=0.0, high=1.0, size=None, dtype=None): high = _torch.as_tensor(high, dtype=dtype, device=device) if bool(_torch.any(low > high)): raise ValueError("Upper bound must be greater than or equal to lower bound") - if size is None: - size = () - elif not hasattr(size, "__iter__"): - size = (size,) - return (high - low) * _torch.rand(size, dtype=dtype, device=device) + low + return (high - low) * _torch.rand(_shape_from_size(size), dtype=dtype, device=device) + low def _tensor_device(*values): @@ -233,11 +245,7 @@ def _floating_distribution_dtype(*values): def _normal_sample_size(size): - if size is None: - return () - if not hasattr(size, "__iter__"): - return (size,) - return tuple(size) + return _shape_from_size(size) @_modify_func_default_dtype(copy=False, kw_only=True) From 392cb5d86103daea5f8f9d48b49f9a58822ccc7b Mon Sep 17 00:00:00 2001 From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com> Date: Mon, 1 Jun 2026 09:46:56 +0200 Subject: [PATCH 2/2] Cover PyTorch random size validation --- tests/backend/test_pytorch_random_backend.py | 55 ++++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 tests/backend/test_pytorch_random_backend.py diff --git a/tests/backend/test_pytorch_random_backend.py b/tests/backend/test_pytorch_random_backend.py new file mode 100644 index 000000000..f364cce20 --- /dev/null +++ b/tests/backend/test_pytorch_random_backend.py @@ -0,0 +1,55 @@ +import pytest + +pytest.importorskip("torch") + +from pyrecest._backend.pytorch import random # noqa: E402 + + +@pytest.mark.parametrize( + "bad_size", + [True, False, (True,), [False, 2], 1.5, (2.0,), "3"], +) +def test_size_arguments_reject_bool_and_non_integral_dimensions(bad_size): + samplers = ( + lambda size: random.rand(size=size), + lambda size: random.uniform(size=size), + lambda size: random.normal(size=size), + lambda size: random.randint(0, 3, size=size), + lambda size: random.choice(3, size=size), + lambda size: random.multivariate_normal([0.0], [[1.0]], size=size), + ) + + for sampler in samplers: + with pytest.raises(TypeError): + sampler(bad_size) + + +@pytest.mark.parametrize("bad_size", [-1, (2, -1)]) +def test_size_arguments_reject_negative_dimensions(bad_size): + samplers = ( + lambda size: random.rand(size=size), + lambda size: random.uniform(size=size), + lambda size: random.normal(size=size), + lambda size: random.randint(0, 3, size=size), + lambda size: random.choice(3, size=size), + lambda size: random.multivariate_normal([0.0], [[1.0]], size=size), + ) + + for sampler in samplers: + with pytest.raises(ValueError): + sampler(bad_size) + + +def test_scalar_and_empty_tuple_sizes_keep_scalar_shape(): + assert random.rand().shape == () + assert random.rand(size=()).shape == () + assert random.normal(size=()).shape == () + assert random.uniform(size=()).shape == () + assert random.randint(0, 3, size=()).shape == () + assert random.multivariate_normal([0.0], [[1.0]], size=()).shape == (1,) + + +def test_zero_sized_choice_still_works_for_empty_population(): + sample = random.choice(0, size=(0,)) + + assert sample.shape == (0,)