"""Stateless PRNG APIs. These are experimental and subject to change without notice. Access via ``torch.func._random``. """ from collections.abc import Sequence import torch def key( seed: int, *, device: torch.device | None = None, impl: str = "philox4x32-10" ) -> torch.Tensor: r"""Create a PRNG key from a seed. A key is a tensor that encodes the state needed to deterministically produce random values. Keys are consumed by generation functions to produce reproducible random tensors without any global state. The internal representation of the key depends on the chosen PRNG algorithm. Args: seed (int): The seed value for the PRNG. device (:class:`torch.device`, optional): The desired device for the returned key. Default: ``cpu``. impl (str): PRNG algorithm. Currently only ``"philox4x32-10"`` is supported. Returns: A tensor representing the PRNG key. .. note:: For the ``"philox4x32-10"`` algorithm, the key is a uint64 tensor of shape ``(2,)`` encoding a ``(seed, offset)`` pair. The offset determines the starting position in the Philox output stream. Example:: >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP """ if impl != "philox4x32-10": raise NotImplementedError(f"key() does not support PRNG impl '{impl}'") # (seed, offset) return torch.tensor([seed, 0], dtype=torch.uint64, device=device) def split(key: torch.Tensor, num: int = 2) -> torch.Tensor: r"""Split a PRNG key into ``num`` new independent keys. Each returned key produces a different, deterministic random sequence. This is the primary mechanism for deriving multiple independent keys from a single parent key without mutating any state. Supports batched keys: if ``key`` has shape ``(*batch, K)``, each key in the batch is split independently and the result has shape ``(num, *batch, K)``. Args: key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or :func:`fold_in`. num (int): Number of keys to produce. Default: ``2``. Returns: A tensor of shape ``(num, *key.shape)`` containing the derived keys. Example:: >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP >>> k1, k2 = torch.func._random.split(key) # doctest: +SKIP """ return torch.ops.aten._philox_key_split(key, num) def fold_in(key: torch.Tensor, data: int | torch.Tensor) -> torch.Tensor: r"""Deterministically derive a new key by folding in an integer value. ``data`` may be a Python ``int`` or a single-item ``uint64`` tensor on the same device as ``key``. Note that passing ``data`` as a tensor prevents it from being baked into a captured CUDA graph, so the graph can be replayed with a different value without recapture. Equivalent to ``split(key, data + 1)[data]``, but more efficient when only a single derived key is needed. Useful for associating a key with a loop iteration, layer index, or other integer identifier. Supports batched keys: if ``key`` has shape ``(*batch, K)``, ``data`` is folded into each key independently. Args: key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or :func:`fold_in`. data (int or Tensor): The value to fold into the key, interpreted as uint64. An ``int`` must be within the inclusive range ``[-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff]``. Negative inputs are remapped to positive values with the formula ``0x1_0000_0000_0000_0000 + data``. A tensor must have dtype ``uint64``, contain a single value, and reside on the same device as ``key``. Returns: A new key tensor with the same shape as ``key``. Example:: >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP >>> k0 = torch.func._random.fold_in(key, 0) # doctest: +SKIP >>> k1 = torch.func._random.fold_in(key, 1) # doctest: +SKIP >>> # Equivalent to split: >>> keys = torch.func._random.split(key, 2) # doctest: +SKIP >>> assert torch.equal(k0, keys[0]) # doctest: +SKIP >>> assert torch.equal(k1, keys[1]) # doctest: +SKIP """ if isinstance(data, torch.Tensor): return torch.ops.aten._philox_key_fold_in.Tensor(key, data) data = int(data) if not -(1 << 63) <= data <= (1 << 64) - 1: raise ValueError( f"fold_in: int data must be in [-2**63, 2**64 - 1], got {data}" ) # Reinterpret as signed int64 due to ATen op schema; kernel will cast back if data >= (1 << 63): data -= 1 << 64 return torch.ops.aten._philox_key_fold_in(key, data) def normal_( key: torch.Tensor, result: torch.Tensor, *, mean: float = 0.0, std: float = 1.0, ) -> torch.Tensor: r"""Fill ``result`` in-place with normal random values from a PRNG key. The values are drawn from a normal distribution with the specified ``mean`` and ``std``. The output is fully determined by the key, so calling with the same key always produces the same result. Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading dimensions of ``result`` must be broadcastable with ``*batch`` and each key independently generates its slice of the output. Args: key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or :func:`fold_in`. result (Tensor): The output tensor to fill in-place. mean (float): Mean of the normal distribution. Default: ``0.0``. std (float): Standard deviation of the normal distribution. Default: ``1.0``. Returns: ``result``, filled with normal random values. Example:: >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP >>> result = torch.empty(1000, device="cuda") # doctest: +SKIP >>> torch.func._random.normal_(key, result) # doctest: +SKIP """ return torch.ops.aten._philox_normal_(result, key, mean, std) def normal( key: torch.Tensor, *shape: tuple[int, ...], mean: float = 0.0, std: float = 1.0, dtype: torch.dtype | None = None, ) -> torch.Tensor: r"""Generate normally distributed random values from a PRNG key. Produces a tensor of the given shape filled with values drawn from a normal distribution with the specified ``mean`` and ``std``. The output is fully determined by the key, so calling with the same key always returns the same result. The output is placed on the same device as ``key``. Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading dimensions of ``shape`` must be broadcastable with ``*batch`` and each key independently generates its slice of the output. Args: key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or :func:`fold_in`. *shape (int): The desired output shape. mean (float): Mean of the normal distribution. Default: ``0.0``. std (float): Standard deviation of the normal distribution. Default: ``1.0``. dtype (:class:`torch.dtype`, optional): The desired dtype. Default: ``torch.float32``. Returns: A tensor of the given shape filled with normal random values. Example:: >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP >>> torch.func._random.normal(key, (1000,)) # doctest: +SKIP """ if len(shape) == 1 and isinstance(shape[0], Sequence): # pyrefly: ignore [bad-argument-type] shape = tuple(shape[0]) if dtype is None: dtype = torch.float32 # pyrefly: ignore [no-matching-overload] result = torch.empty(shape, dtype=dtype, device=key.device) return normal_(key, result, mean=mean, std=std) def uniform_( key: torch.Tensor, result: torch.Tensor, *, low: float = 0.0, high: float = 1.0, ) -> torch.Tensor: r"""Fill ``result`` in-place with uniform random values from a PRNG key. The values are drawn uniformly from the interval ``[low, high)``. The output is fully determined by the key, so calling with the same key always produces the same result. Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading dimensions of ``result`` must be broadcastable with ``*batch`` and each key independently generates its slice of the output. Args: key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or :func:`fold_in`. result (Tensor): The output tensor to fill in-place. low (float): Lower bound (inclusive) of the uniform distribution. Default: ``0.0``. high (float): Upper bound (exclusive) of the uniform distribution. Default: ``1.0``. Returns: ``result``, filled with uniform random values. Example:: >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP >>> result = torch.empty(1000, device="cuda") # doctest: +SKIP >>> torch.func._random.uniform_(key, result) # doctest: +SKIP """ return torch.ops.aten._philox_uniform_(result, key, low, high) def uniform( key: torch.Tensor, *shape: tuple[int, ...], low: float = 0.0, high: float = 1.0, dtype: torch.dtype | None = None, ) -> torch.Tensor: r"""Generate uniformly distributed random values from a PRNG key. Produces a tensor of the given shape filled with values drawn uniformly from the interval ``[low, high)``. The output is fully determined by the key, so calling with the same key always returns the same result. The output is placed on the same device as ``key``. Supports batched keys: if ``key`` has shape ``(*batch, K)``, the leading dimensions of ``shape`` must be broadcastable with ``*batch`` and each key independently generates its slice of the output. Args: key (Tensor): A PRNG key returned by :func:`key`, :func:`split`, or :func:`fold_in`. *shape (int): The desired output shape. low (float): Lower bound (inclusive) of the uniform distribution. Default: ``0.0``. high (float): Upper bound (exclusive) of the uniform distribution. Default: ``1.0``. dtype (:class:`torch.dtype`, optional): The desired dtype. Default: ``torch.float32``. Returns: A tensor of the given shape filled with uniform random values. Example:: >>> key = torch.func._random.key(42, device="cuda") # doctest: +SKIP >>> torch.func._random.uniform(key, (1000,)) # doctest: +SKIP """ if len(shape) == 1 and isinstance(shape[0], Sequence): # pyrefly: ignore [bad-argument-type] shape = tuple(shape[0]) if dtype is None: dtype = torch.float32 # pyrefly: ignore [no-matching-overload] result = torch.empty(shape, dtype=dtype, device=key.device) return uniform_(key, result, low=low, high=high)