diff --git a/pyproject.toml b/pyproject.toml index a23865e381..c6c7bc3f48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -147,6 +147,7 @@ unfixable = [ [tool.ruff.lint.isort] lines-after-imports = 2 +known-first-party = ["pytensor", "tests"] [tool.ruff.lint.per-file-ignores] # TODO: Get rid of these: diff --git a/pytensor/link/mlx/dispatch/__init__.py b/pytensor/link/mlx/dispatch/__init__.py index 59b0604856..222191415c 100644 --- a/pytensor/link/mlx/dispatch/__init__.py +++ b/pytensor/link/mlx/dispatch/__init__.py @@ -17,4 +17,5 @@ import pytensor.link.mlx.dispatch.pad import pytensor.link.mlx.dispatch.sort import pytensor.link.mlx.dispatch.linalg +import pytensor.link.mlx.dispatch.random # isort: on diff --git a/pytensor/link/mlx/dispatch/random.py b/pytensor/link/mlx/dispatch/random.py new file mode 100644 index 0000000000..ab8d9d71be --- /dev/null +++ b/pytensor/link/mlx/dispatch/random.py @@ -0,0 +1,276 @@ +from functools import singledispatch + +import mlx.core as mx +from numpy.random import Generator + +import pytensor.tensor.random.basic as ptr +from pytensor.link.mlx.dispatch.basic import mlx_funcify, mlx_typify +from pytensor.link.mlx.dispatch.tensor_basic import ( + convert_dtype_to_mlx, + mlx_to_list_shape, +) + + +def numpy_generator_to_mlx_key(rng: Generator) -> mx.array: + """Convert a NumPy Generator to an MLX random key. + + MLX keys are 64-bit, so we XOR-fold the two halves of the 128-bit PCG64 + state to keep all of its entropy. + """ + state = rng.bit_generator.state + if state["bit_generator"] not in ("PCG64", "PCG64DXSM"): + raise NotImplementedError( + "MLX RNG conversion only supports the PCG64 bit generator, got " + f"{state['bit_generator']}." + ) + state_128 = int(state["state"]["state"]) + upper = (state_128 >> 64) & 0xFFFFFFFFFFFFFFFF + lower = state_128 & 0xFFFFFFFFFFFFFFFF + return mx.random.key(upper ^ lower) + + +def _shape_from_size(size, *parameters) -> list[int] | tuple[int, ...]: + """Sampling shape: ``size`` if given, else the broadcast of the parameters.""" + if size is not None: + return mlx_to_list_shape(size) + return tuple(mx.broadcast_shapes(*(p.shape for p in parameters))) + + +@mlx_typify.register(Generator) +def mlx_typify_Generator(rng, **kwargs): + return numpy_generator_to_mlx_key(rng) + + +@mlx_funcify.register(ptr.RandomVariable) +def mlx_funcify_RandomVariable(op, node, **kwargs): + rv = node.outputs[1] + out_dtype = rv.type.dtype + + # MLX random primitives reject float64 regardless of device, so sample in + # float32 and widen the draw to the declared output dtype afterwards. + mlx_out_dtype = convert_dtype_to_mlx(out_dtype) + sample_dtype = mx.float32 if mlx_out_dtype == mx.float64 else mlx_out_dtype + + sample_fn_inner = mlx_sample_fn(op, node) + + def sample_fn(rng, size, *parameters): + new_rng, sampling_key = mx.random.split(rng, num=2) + sample = sample_fn_inner(sampling_key, size, sample_dtype, *parameters) + if sample.dtype != mlx_out_dtype: + sample = sample.astype(mlx_out_dtype) + return (new_rng, sample) + + return sample_fn + + +@singledispatch +def mlx_sample_fn(op, node): + raise NotImplementedError( + f"No MLX implementation for the given distribution: {op.name}" + ) + + +@mlx_sample_fn.register(ptr.NormalRV) +def mlx_sample_fn_normal(op, node): + def sample_fn(rng_key, size, dtype, loc, scale): + mlx_dtype = convert_dtype_to_mlx(dtype) + loc = mx.array(loc, dtype=mlx_dtype) + scale = mx.array(scale, dtype=mlx_dtype) + shape = _shape_from_size(size, loc, scale) + return loc + scale * mx.random.normal(shape=shape, dtype=mlx_dtype, key=rng_key) + + return sample_fn + + +@mlx_sample_fn.register(ptr.LaplaceRV) +def mlx_sample_fn_laplace(op, node): + def sample_fn(rng_key, size, dtype, loc, scale): + mlx_dtype = convert_dtype_to_mlx(dtype) + loc = mx.array(loc, dtype=mlx_dtype) + scale = mx.array(scale, dtype=mlx_dtype) + shape = _shape_from_size(size, loc, scale) + return loc + scale * mx.random.laplace( + shape=shape, dtype=mlx_dtype, key=rng_key + ) + + return sample_fn + + +@mlx_sample_fn.register(ptr.GumbelRV) +def mlx_sample_fn_gumbel(op, node): + def sample_fn(rng_key, size, dtype, loc, scale): + mlx_dtype = convert_dtype_to_mlx(dtype) + loc = mx.array(loc, dtype=mlx_dtype) + scale = mx.array(scale, dtype=mlx_dtype) + shape = _shape_from_size(size, loc, scale) + return loc + scale * mx.random.gumbel(shape=shape, dtype=mlx_dtype, key=rng_key) + + return sample_fn + + +@mlx_sample_fn.register(ptr.UniformRV) +def mlx_sample_fn_uniform(op, node): + def sample_fn(rng_key, size, dtype, low, high): + mlx_dtype = convert_dtype_to_mlx(dtype) + low = mx.array(low, dtype=mlx_dtype) + high = mx.array(high, dtype=mlx_dtype) + shape = _shape_from_size(size, low, high) + return mx.random.uniform( + low=low, high=high, shape=shape, dtype=mlx_dtype, key=rng_key + ) + + return sample_fn + + +@mlx_sample_fn.register(ptr.IntegersRV) +def mlx_sample_fn_integers(op, node): + def sample_fn(rng_key, size, dtype, low, high): + low = mx.array(low) + high = mx.array(high) + shape = _shape_from_size(size, low, high) + # Sample at full int64 width and cast the result: PyTensor casts the + # output, not the bounds, so narrow/wide dtypes don't corrupt the range. + return mx.random.randint( + low=low, high=high, shape=shape, dtype=mx.int64, key=rng_key + ).astype(convert_dtype_to_mlx(dtype)) + + return sample_fn + + +@mlx_sample_fn.register(ptr.BernoulliRV) +def mlx_sample_fn_bernoulli(op, node): + def sample_fn(rng_key, size, dtype, p): + p = mx.array(p) + shape = mlx_to_list_shape(size) if size is not None else None + # MLX draws bool; PyTensor declares an int dtype. + return mx.random.bernoulli(p=p, shape=shape, key=rng_key).astype( + convert_dtype_to_mlx(dtype) + ) + + return sample_fn + + +@mlx_sample_fn.register(ptr.CategoricalRV) +def mlx_sample_fn_categorical(op, node): + def sample_fn(rng_key, size, dtype, p): + logits = mx.log(mx.array(p)) + shape = mlx_to_list_shape(size) if size is not None else None + # MLX draws uint32; PyTensor declares an int dtype. + return mx.random.categorical( + logits=logits, axis=-1, shape=shape, key=rng_key + ).astype(convert_dtype_to_mlx(dtype)) + + return sample_fn + + +@mlx_sample_fn.register(ptr.MvNormalRV) +def mlx_sample_fn_mvnormal(op, node): + method = op.method + + def sample_fn(rng_key, size, dtype, mean, cov): + mlx_dtype = convert_dtype_to_mlx(dtype) + mean = mx.array(mean, dtype=mlx_dtype) + cov = mx.array(cov, dtype=mlx_dtype) + + n = cov.shape[-1] + if size is not None: + batch_shape = mlx_to_list_shape(size) + else: + batch_shape = mx.broadcast_shapes(mean.shape[:-1], cov.shape[:-2]) + + if 0 in tuple(batch_shape): + # Empty batch dim crashes MLX's compiled matmul; the draw is empty anyway. + return mx.broadcast_to(mean, [*batch_shape, n]) + + # Factor ``cov = A @ A.T`` so that ``mean + A @ z`` has covariance ``cov``. + if method == "cholesky": + A = mx.linalg.cholesky(cov, stream=mx.cpu) + elif method == "svd": + U, s, _ = mx.linalg.svd(cov, stream=mx.cpu) + A = U * mx.sqrt(s)[..., None, :] + else: # eigh + w, vecs = mx.linalg.eigh(cov, stream=mx.cpu) + A = vecs * mx.sqrt(w)[..., None, :] + + z = mx.random.normal(shape=[*batch_shape, n], dtype=mlx_dtype, key=rng_key) + return mean + (A @ z[..., None])[..., 0] + + return sample_fn + + +@mlx_sample_fn.register(ptr.PermutationRV) +def mlx_sample_fn_permutation(op, node): + if op.batch_ndim(node): + raise NotImplementedError( + "MLX random.permutation does not support batch dimensions." + ) + + def sample_fn(rng_key, size, dtype, x): + return mx.random.permutation(x, key=rng_key) + + return sample_fn + + +@mlx_sample_fn.register(ptr.LogNormalRV) +def mlx_sample_fn_lognormal(op, node): + def sample_fn(rng_key, size, dtype, mu, sigma): + mlx_dtype = convert_dtype_to_mlx(dtype) + mu = mx.array(mu, dtype=mlx_dtype) + sigma = mx.array(sigma, dtype=mlx_dtype) + shape = _shape_from_size(size, mu, sigma) + z = mx.random.normal(shape=shape, dtype=mlx_dtype, key=rng_key) + return mx.exp(mu + sigma * z) + + return sample_fn + + +@mlx_sample_fn.register(ptr.HalfNormalRV) +def mlx_sample_fn_halfnormal(op, node): + def sample_fn(rng_key, size, dtype, loc, scale): + mlx_dtype = convert_dtype_to_mlx(dtype) + loc = mx.array(loc, dtype=mlx_dtype) + scale = mx.array(scale, dtype=mlx_dtype) + shape = _shape_from_size(size, loc, scale) + z = mx.random.normal(shape=shape, dtype=mlx_dtype, key=rng_key) + return loc + scale * mx.abs(z) + + return sample_fn + + +@mlx_sample_fn.register(ptr.ExponentialRV) +def mlx_sample_fn_exponential(op, node): + def sample_fn(rng_key, size, dtype, scale): + mlx_dtype = convert_dtype_to_mlx(dtype) + scale = mx.array(scale, dtype=mlx_dtype) + shape = _shape_from_size(size, scale) + u = mx.random.uniform(shape=shape, dtype=mlx_dtype, key=rng_key) + # log1p(-u) avoids -inf when u=0 (u is in [0, 1)) + return -scale * mx.log1p(-u) + + return sample_fn + + +@mlx_sample_fn.register(ptr.LogisticRV) +def mlx_sample_fn_logistic(op, node): + def sample_fn(rng_key, size, dtype, loc, scale): + mlx_dtype = convert_dtype_to_mlx(dtype) + loc = mx.array(loc, dtype=mlx_dtype) + scale = mx.array(scale, dtype=mlx_dtype) + shape = _shape_from_size(size, loc, scale) + u = mx.random.uniform(shape=shape, dtype=mlx_dtype, key=rng_key) + return loc + scale * mx.log(u / (1 - u)) + + return sample_fn + + +@mlx_sample_fn.register(ptr.CauchyRV) +def mlx_sample_fn_cauchy(op, node): + def sample_fn(rng_key, size, dtype, loc, scale): + mlx_dtype = convert_dtype_to_mlx(dtype) + loc = mx.array(loc, dtype=mlx_dtype) + scale = mx.array(scale, dtype=mlx_dtype) + shape = _shape_from_size(size, loc, scale) + u = mx.random.uniform(shape=shape, dtype=mlx_dtype, key=rng_key) + return loc + scale * mx.tan(mx.pi * (u - 0.5)) + + return sample_fn diff --git a/pytensor/link/mlx/dispatch/tensor_basic.py b/pytensor/link/mlx/dispatch/tensor_basic.py index 3cdc47323f..b4ae841eac 100644 --- a/pytensor/link/mlx/dispatch/tensor_basic.py +++ b/pytensor/link/mlx/dispatch/tensor_basic.py @@ -239,6 +239,16 @@ def _coerce_to_int(value): raise +def mlx_to_list_shape(size) -> list[int]: + """Convert a size value (mx.array, np.ndarray, or sequence) to a plain Python list of ints. + + Used by random variable dispatch to normalise the ``size`` argument, which + PyTensor may pass as an ``mx.array`` or ``np.ndarray`` rather than a plain + Python list. + """ + return [_coerce_to_int(x) for x in size] + + def _rethrow_dynamic_shape_error(exc): msg = str(exc) if "[eval] Attempting to eval an array during function transformations" in msg: diff --git a/pytensor/link/mlx/linker.py b/pytensor/link/mlx/linker.py index 9d662308cf..b8d44f1c0b 100644 --- a/pytensor/link/mlx/linker.py +++ b/pytensor/link/mlx/linker.py @@ -1,3 +1,6 @@ +import warnings + +from pytensor.compile.sharedvalue import SharedVariable, shared from pytensor.link.basic import JITLinker @@ -19,7 +22,7 @@ def __init__(self, use_compile=True, *args, **kwargs): self.gen_functors = [] self.use_compile = use_compile - def fgraph_convert(self, fgraph, **kwargs): + def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs): """Convert a PyTensor FunctionGraph to an MLX-compatible function. Parameters @@ -33,9 +36,63 @@ def fgraph_convert(self, fgraph, **kwargs): An MLX-compatible function """ from pytensor.link.mlx.dispatch import mlx_funcify + from pytensor.tensor.random.type import RandomType + + shared_rng_inputs = [ + inp + for inp in fgraph.inputs + if (isinstance(inp, SharedVariable) and isinstance(inp.type, RandomType)) + ] + + # Replace any shared RNG inputs so that their values can be updated in place + # without affecting the original RNG container. This is necessary because + # MLX does not accept Generators as inputs, and they will have to + # be typified + if shared_rng_inputs: + warnings.warn( + f"The RandomType SharedVariables {shared_rng_inputs} will not be used " + f"in the compiled MLX graph. Instead a copy will be used.", + UserWarning, + ) + new_shared_rng_inputs = [ + shared(inp.get_value(borrow=False)) for inp in shared_rng_inputs + ] + + fgraph.replace_all( + zip(shared_rng_inputs, new_shared_rng_inputs, strict=True), + import_missing=True, + reason="MLXLinker.fgraph_convert", + ) + + for old_inp, new_inp in zip( + shared_rng_inputs, new_shared_rng_inputs, strict=True + ): + new_inp_storage = [new_inp.get_value(borrow=True)] + storage_map[new_inp] = new_inp_storage + old_inp_storage = storage_map.pop(old_inp) + # Find index of old_inp_storage in input_storage + for input_storage_idx, input_storage_item in enumerate(input_storage): + # We have to establish equality based on identity because input_storage may contain numpy arrays + if input_storage_item is old_inp_storage: + break + else: # no break + raise ValueError() + input_storage[input_storage_idx] = new_inp_storage + # We need to change the order of the inputs of the FunctionGraph + # so that the new input is in the same position as to old one, + # to align with the storage_map. We hope this is safe! + old_inp_fgraph_index = fgraph.inputs.index(old_inp) + fgraph.remove_input( + old_inp_fgraph_index, + reason="MLXLinker.fgraph_convert", + ) + fgraph.inputs.remove(new_inp) + fgraph.inputs.insert(old_inp_fgraph_index, new_inp) return mlx_funcify( fgraph, + input_storage=input_storage, + storage_map=storage_map, **kwargs, ) @@ -71,9 +128,16 @@ def create_thunk_inputs(self, storage_map): list The inputs for the thunk """ + from numpy.random import Generator + + from pytensor.link.mlx.dispatch import mlx_typify + thunk_inputs = [] for n in self.fgraph.inputs: sinput = storage_map[n] + if isinstance(sinput[0], Generator): + # Convert Generator into MLX PRNG key + sinput[0] = mlx_typify(sinput[0]) thunk_inputs.append(sinput) return thunk_inputs diff --git a/tests/link/mlx/test_random.py b/tests/link/mlx/test_random.py new file mode 100644 index 0000000000..073c6f2818 --- /dev/null +++ b/tests/link/mlx/test_random.py @@ -0,0 +1,316 @@ +import numpy as np +import pytest + +import pytensor +import pytensor.tensor as pt +from pytensor.compile.maker import function +from pytensor.compile.sharedvalue import shared +from pytensor.tensor.random.utils import RandomStream + + +mx = pytest.importorskip("mlx.core") + + +def test_normal_cumsum(): + out = pt.random.normal(size=(52,)).cumsum() + result = out.eval(mode="MLX") + assert isinstance(result, mx.array) + assert result.shape == (52,) + + +def check_shape_and_dtype( + make_rv, expected_shape, expected_dtype=None, n_evals=2, mode="MLX" +): + """Compile and run an RV under MLX; assert shape/dtype and that successive + draws differ (RNG state is threaded).""" + srng = RandomStream(seed=12345) + rv = make_rv(srng) + f = pytensor.function([], rv, mode=mode, updates=srng.updates()) + results = [np.array(f()) for _ in range(n_evals)] + + for r in results: + assert r.shape == expected_shape, ( + f"Expected shape {expected_shape}, got {r.shape}" + ) + if expected_dtype is not None: + if expected_dtype == "float": + assert np.issubdtype(r.dtype, np.floating), ( + f"Expected a float dtype, got {r.dtype}" + ) + else: + assert r.dtype == np.dtype(expected_dtype), ( + f"Expected dtype {expected_dtype}, got {r.dtype}" + ) + + assert not np.array_equal(results[0], results[1]), ( + "Two draws were identical — RNG not advancing" + ) + + return results + + +@pytest.mark.parametrize( + "make_rv,expected_shape,expected_dtype", + [ + (lambda srng: srng.normal(loc=0.0, scale=1.0, size=(3, 4)), (3, 4), "float"), + (lambda srng: srng.normal(loc=2.0, scale=0.5), (), None), + (lambda srng: srng.uniform(low=0.0, high=1.0, size=(10,)), (10,), "float"), + (lambda srng: srng.bernoulli(p=0.7, size=(5, 5)), (5, 5), "int64"), + ( + lambda srng: srng.categorical( + p=np.array([0.1, 0.4, 0.5], dtype="float32"), size=(8,) + ), + (8,), + "int64", + ), + (lambda srng: srng.laplace(loc=0.0, scale=1.0, size=(7,)), (7,), "float"), + (lambda srng: srng.gumbel(loc=0.0, scale=1.0, size=(6,)), (6,), "float"), + (lambda srng: srng.lognormal(mu=0.0, sigma=1.0, size=(5,)), (5,), "float"), + (lambda srng: srng.lognormal(mu=0.0, sigma=1.0), (), None), + (lambda srng: srng.halfnormal(loc=0.0, scale=1.0, size=(4,)), (4,), "float"), + (lambda srng: srng.halfnormal(loc=0.0, scale=1.0), (), None), + (lambda srng: srng.exponential(scale=1.0, size=(6,)), (6,), "float"), + (lambda srng: srng.logistic(loc=0.0, scale=1.0, size=(7,)), (7,), "float"), + (lambda srng: srng.cauchy(loc=0.0, scale=1.0, size=(8,)), (8,), "float"), + ], +) +def test_shape_dtype(make_rv, expected_shape, expected_dtype): + check_shape_and_dtype(make_rv, expected_shape, expected_dtype) + + +def test_normal_array_params(): + result = pt.random.normal(loc=[0, 1], scale=[1.0, 0.3], size=(100, 2)).eval( + mode="MLX" + ) + assert result.shape == (100, 2) + means = np.array(result).mean(axis=0) + assert abs(means[0]) < 0.3 + assert abs(means[1] - 1.0) < 0.3 + + +def test_uniform_values(): + results = check_shape_and_dtype( + lambda srng: srng.uniform(low=0.0, high=1.0, size=(10,)), + (10,), + "float", + ) + r = np.array(results[0]) + assert np.all(r >= 0.0) + assert np.all(r < 1.0) + + +def test_categorical_values(): + probs = np.array([0.1, 0.4, 0.5], dtype=np.float32) + results = check_shape_and_dtype( + lambda srng: srng.categorical(p=probs, size=(8,)), + (8,), + "int64", + ) + r = np.array(results[0]) + assert np.all(r < 3) + assert np.all(r >= 0) + + +def test_mvnormal_shape(): + mean = np.zeros(4, dtype=np.float32) + cov = np.eye(4, dtype=np.float32) + check_shape_and_dtype( + lambda srng: srng.multivariate_normal(mean=mean, cov=cov, size=(6,)), + (6, 4), + "float", + ) + + +@pytest.mark.parametrize("method", ["cholesky", "svd", "eigh"]) +def test_mvnormal_decomposition_method(method): + mean = np.zeros(4, dtype=np.float32) + cov = np.eye(4, dtype=np.float32) + check_shape_and_dtype( + lambda srng: srng.multivariate_normal( + mean=mean, cov=cov, size=(6,), method=method + ), + (6, 4), + "float", + ) + + +def test_mvnormal_batched_params_with_size(): + # Batched covariances combined with an explicit ``size`` must broadcast + # rather than reshape a single matrix (regression for a reshape crash). + mean = np.zeros((2, 3), dtype=np.float32) + cov = np.stack([np.eye(3) * 0.01, np.eye(3) * 9.0]).astype(np.float32) + check_shape_and_dtype( + lambda srng: srng.multivariate_normal(mean=mean, cov=cov, size=(2,)), + (2, 3), + "float", + ) + + +def test_mvnormal_empty_batch(): + # An empty batch dim used to segfault the MLX compiled matmul path; it must + # return an empty array of the broadcast output shape instead. + mean = np.zeros(3, dtype=np.float32) + cov = np.eye(3, dtype=np.float32) + result = pt.random.multivariate_normal(mean=mean, cov=cov, size=(0,)).eval( + mode="MLX" + ) + assert np.array(result).shape == (0, 3) + + +def test_integers_shape(): + results = check_shape_and_dtype( + lambda srng: srng.integers(low=0, high=10, size=(12,)), + (12,), + ) + r = np.array(results[0]) + assert np.all(r >= 0) + assert np.all(r < 10) + + +def test_integers_narrow_dtype_wraps(): + # Sampling [250, 300) then casting to uint8 must wrap (regression for + # casting the bounds first, which collapsed the interval to one value). + r = np.array( + pt.random.integers(low=250, high=300, size=(20_000,), dtype="uint8").eval( + mode="MLX" + ) + ) + assert r.dtype == np.uint8 + assert len(np.unique(r)) > 1 + + +def test_integers_wide_bounds(): + # Default int64 draws must sample at full width, not MLX's default int32 + # (regression for bounds above 2**31 piling up at the int32 max). + r = np.array( + pt.random.integers(low=0, high=3_000_000_000, size=(20_000,)).eval(mode="MLX") + ) + assert r.max() > 2**31 + + +def test_permutation_shape(): + x = np.arange(8, dtype=np.int32) + results = check_shape_and_dtype( + lambda srng: srng.permutation(x), + (8,), + ) + assert sorted(np.array(results[0]).tolist()) == list(range(8)) + + +def test_lognormal_shape_dtype(): + results = check_shape_and_dtype( + lambda srng: srng.lognormal(mu=0.0, sigma=1.0, size=(5,)), + (5,), + "float", + ) + r = np.array(results[0]) + assert np.all(r > 0) + + +def test_halfnormal_nonzero_loc(): + # HalfNormal is ``loc + scale * |z|`` (support ``[loc, inf)``), not + # ``|loc + scale * z|``. Draw a large sample and check the support bound. + loc, scale = 5.0, 2.0 + r = np.array( + pt.random.halfnormal(loc=loc, scale=scale, size=(100_000,)).eval(mode="MLX") + ) + assert r.min() >= loc - 1e-4 + assert abs(r.mean() - (loc + scale * np.sqrt(2 / np.pi))) < 0.1 + + +def test_exponential_shape_dtype(): + results = check_shape_and_dtype( + lambda srng: srng.exponential(scale=1.0, size=(6,)), + (6,), + "float", + ) + r = np.array(results[0]) + assert np.all(r > 0) + + +def test_non_pcg64_generator_raises(): + # Only PCG64 state can be folded into an MLX key; other bit generators must + # fail loudly rather than with an opaque KeyError. + from pytensor.link.mlx.dispatch.random import numpy_generator_to_mlx_key + + with pytest.raises(NotImplementedError, match="PCG64"): + numpy_generator_to_mlx_key(np.random.Generator(np.random.MT19937(0))) + + +def compile_shared_rng_function(*args, mode="MLX", **kwargs): + with pytest.warns( + UserWarning, match=r"The RandomType SharedVariables \[.+\] will not be used" + ): + return function(*args, mode=mode, **kwargs) + + +def test_random_updates(): + original_value = np.random.default_rng(seed=98) + rng = shared(original_value, name="original_rng", borrow=False) + next_rng, x = pt.random.normal(name="x", rng=rng).owner.outputs + + f = compile_shared_rng_function([], [x], updates={rng: next_rng}) + assert f() != f() + + # Check that the original shared variable was not overwritten when typifying + assert all( + a == b if not isinstance(a, np.ndarray) else np.array_equal(a, b) + for a, b in zip( + rng.get_value().bit_generator.state, + original_value.bit_generator.state, + strict=True, + ) + ) + + +@pytest.mark.parametrize("noise_first", (False, True)) +def test_replaced_shared_rng_storage_order(noise_first): + # Test that replacing the RNG variable in the linker does not cause + # a disalignment between the compiled graph and the storage_map. + + mu = pytensor.shared(np.array(1.0), name="mu") + rng = pytensor.shared(np.random.default_rng(123)) + next_rng, noise = pt.random.normal(rng=rng).owner.outputs + + out = noise * mu if noise_first else mu * noise + + updates = { + mu: pt.grad(out, mu), + rng: next_rng, + } + f = compile_shared_rng_function([], [out], updates=updates) + + # Confirm that input_storage type and fgraph input order are aligned + for storage, fgraph_input in zip( + f.input_storage, f.maker.fgraph.inputs, strict=True + ): + assert storage.type == fgraph_input.type + + assert mu.get_value() == 1 + f() + assert mu.get_value() != 1 + + +def test_replaced_shared_rng_storage_ordering_equality(): + """Test that storage identity comparison works when numpy arrays precede + the RNG in input_storage (regression test for issue #314).""" + pt_rng = RandomStream(1) + + batchshape = (3, 1, 4, 4) + inp_shared = pytensor.shared( + np.zeros(batchshape, dtype="float64"), name="inp_shared" + ) + + inp = pt.tensor4(dtype="float64", name="inp") + inp_update = inp + pt_rng.normal(size=inp.shape, loc=5, scale=1e-5) + + fn = compile_shared_rng_function( + inputs=[], + outputs=[], + updates={inp_shared: inp_update}, + givens={inp: inp_shared}, + ) + fn() + np.testing.assert_allclose(np.array(inp_shared.get_value()), 5, rtol=1e-2) + fn() + np.testing.assert_allclose(np.array(inp_shared.get_value()), 10, rtol=1e-2)