orng 0.1.0__py3-none-any.whl

orng/__init__.py

@@ -0,0 +1,23 @@
+import importlib.metadata
+
+try:
+    __version__ = importlib.metadata.version("orng")
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "unknown"
+
+from .backends import create_backend_from_xp, infer_backend_name_from_xp
+from .functional import (
+    create_functional_backend,
+    create_functional_backend_from_xp,
+)
+from .orng import ArrayRNG, RandomGenerator
+
+__all__ = [
+    "ArrayRNG",
+    "RandomGenerator",
+    "create_backend_from_xp",
+    "create_functional_backend",
+    "create_functional_backend_from_xp",
+    "infer_backend_name_from_xp",
+    "__version__",
+]

orng/_utils.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import math
+from typing import Sequence, Tuple
+
+SizeLike = int | Sequence[int] | None
+
+
+def normalize_shape(size: SizeLike) -> Tuple[int, ...]:
+    """Convert a size argument to a canonical ``tuple`` form."""
+    if size is None:
+        return ()
+    if isinstance(size, int):
+        if size < 0:
+            raise ValueError("size must be non-negative.")
+        return (size,)
+    shape = tuple(int(dim) for dim in size)
+    if any(dim < 0 for dim in shape):
+        raise ValueError("size entries must be non-negative.")
+    return shape
+
+
+def total_size(shape: Tuple[int, ...]) -> int:
+    if not shape:
+        return 1
+    return math.prod(shape)
+
+
+__all__ = ["SizeLike", "normalize_shape", "total_size"]

orng/backends/__init__.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Dict
+
+from .cupy import CuPyBackend
+from .jax import JAXBackend
+from .numpy import NumPyBackend
+from .torch import TorchBackend
+
+BackendFactory = Callable[..., Any]
+
+_FACTORIES: Dict[str, BackendFactory] = {
+    "numpy": lambda *, seed=None, generator=None, device=None: NumPyBackend(
+        seed=seed,
+        generator=generator,
+    ),
+    "cupy": lambda *, seed=None, generator=None, device=None: CuPyBackend(
+        seed=seed,
+        generator=generator,
+    ),
+    "torch": lambda *, seed=None, generator=None, device=None: TorchBackend(
+        seed=seed,
+        generator=generator,
+        device=device,
+    ),
+    "pytorch": lambda *, seed=None, generator=None, device=None: TorchBackend(
+        seed=seed,
+        generator=generator,
+        device=device,
+    ),
+    "jax": lambda *, seed=None, generator=None, device=None: JAXBackend(
+        seed=seed,
+        key=generator,
+    ),
+}
+
+
+def infer_backend_name_from_xp(xp: Any) -> str:
+    try:
+        from array_api_compat import (
+            is_cupy_namespace,
+            is_jax_namespace,
+            is_numpy_namespace,
+            is_torch_namespace,
+        )
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        raise ImportError(
+            "Inferring ORNG backends from an array namespace requires "
+            "'array_api_compat'."
+        ) from exc
+
+    if is_numpy_namespace(xp):
+        return "numpy"
+    if is_jax_namespace(xp):
+        return "jax"
+    if is_torch_namespace(xp):
+        return "torch"
+    if is_cupy_namespace(xp):
+        return "cupy"
+    raise ValueError("Unsupported array namespace for ORNG backend inference.")
+
+
+def create_backend(
+    name: str,
+    *,
+    seed: int | None,
+    generator: Any | None,
+    device: Any | None,
+):
+    try:
+        factory = _FACTORIES[name.lower()]
+    except KeyError as exc:  # pragma: no cover - defensive
+        supported = "', '".join(
+            sorted({k for k in _FACTORIES if k != "pytorch"})
+        )
+        raise ValueError(
+            f"Unsupported backend '{name}'. Expected one of '{supported}'."
+        ) from exc
+    return factory(seed=seed, generator=generator, device=device)
+
+
+def create_backend_from_xp(
+    xp: Any,
+    *,
+    seed: int | None,
+    generator: Any | None,
+    device: Any | None,
+):
+    return create_backend(
+        infer_backend_name_from_xp(xp),
+        seed=seed,
+        generator=generator,
+        device=device,
+    )
+
+
+__all__ = [
+    "create_backend",
+    "create_backend_from_xp",
+    "infer_backend_name_from_xp",
+]

orng/backends/cupy.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+import copy
+from typing import Any
+
+from .._utils import SizeLike
+
+
+class CuPyBackend:
+    def __init__(self, *, seed: int | None, generator: Any | None) -> None:
+        self._impl = CuPyFunctionalBackend(pure=False)
+        self._state = self._impl.init_state(seed=seed, generator=generator)
+
+    def random(self, *, size: SizeLike, dtype: Any | None) -> Any:
+        result, self._state = self._impl.random(
+            self._state,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def uniform(
+        self,
+        *,
+        low: Any,
+        high: Any,
+        size: SizeLike,
+        dtype: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.uniform(
+            self._state,
+            low=low,
+            high=high,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def normal(
+        self,
+        *,
+        loc: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.normal(
+            self._state,
+            loc=loc,
+            scale=scale,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def gamma(
+        self,
+        *,
+        shape: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.gamma(
+            self._state,
+            shape=shape,
+            scale=scale,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def choice(
+        self,
+        population: int | Any,
+        *,
+        size: SizeLike,
+        replace: bool,
+        probabilities: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.choice(
+            self._state,
+            population,
+            size=size,
+            replace=replace,
+            probabilities=probabilities,
+        )
+        return result
+
+
+class CuPyFunctionalBackend:
+    def __init__(self, *, pure: bool = True) -> None:
+        try:
+            import cupy as cp
+        except ImportError as exc:  # pragma: no cover - optional dependency
+            raise ImportError(
+                "CuPy backend requires the 'cupy' package to be installed. "
+                "Install it with `pip install orng[cupy]`."
+            ) from exc
+        self._cupy = cp
+        self._pure = pure
+
+    def init_state(self, *, seed: int | None, generator: Any | None) -> Any:
+        cp = self._cupy
+        if generator is None:
+            gen = cp.random.default_rng(seed)
+        elif isinstance(generator, cp.random.Generator):
+            gen = generator
+        else:
+            raise TypeError(
+                "generator must be a cupy.random.Generator when using the "
+                "CuPy backend."
+            )
+        if self._pure:
+            return copy.deepcopy(gen.bit_generator.state)
+        return gen
+
+    def _generator_from_state(self, state: Any) -> Any:
+        if not self._pure:
+            if not isinstance(state, self._cupy.random.Generator):
+                raise TypeError(
+                    "state must be a cupy.random.Generator when pure=False."
+                )
+            return state
+        gen = self._cupy.random.default_rng()
+        gen.bit_generator.state = copy.deepcopy(state)
+        return gen
+
+    def _result_and_next_state(self, gen: Any, result: Any) -> tuple[Any, Any]:
+        if self._pure:
+            return result, copy.deepcopy(gen.bit_generator.state)
+        return result, gen
+
+    def random(
+        self,
+        state: Any,
+        *,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        gen = self._generator_from_state(state)
+        result = gen.random(size=size, dtype=dtype)
+        return self._result_and_next_state(gen, result)
+
+    def uniform(
+        self,
+        state: Any,
+        *,
+        low: Any,
+        high: Any,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        gen = self._generator_from_state(state)
+        result = gen.uniform(low=low, high=high, size=size, dtype=dtype)
+        return self._result_and_next_state(gen, result)
+
+    def normal(
+        self,
+        state: Any,
+        *,
+        loc: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        gen = self._generator_from_state(state)
+        standard = gen.standard_normal(size=size, dtype=dtype)
+        result = loc + standard * scale
+        return self._result_and_next_state(gen, result)
+
+    def gamma(
+        self,
+        state: Any,
+        *,
+        shape: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        gen = self._generator_from_state(state)
+        result = gen.gamma(shape=shape, scale=scale, size=size)
+        if dtype is not None:
+            result = self._cupy.asarray(result, dtype=dtype)
+        return self._result_and_next_state(gen, result)
+
+    def choice(
+        self,
+        state: Any,
+        population: int | Any,
+        *,
+        size: SizeLike,
+        replace: bool,
+        probabilities: Any | None,
+    ) -> tuple[Any, Any]:
+        cp = self._cupy
+        gen = self._generator_from_state(state)
+
+        if isinstance(population, int):
+            n = population
+            values = None
+        else:
+            values = cp.asarray(population)
+            n = values.shape[0]
+
+        target_size = size if size is not None else ()
+
+        probs = None
+        if probabilities is not None:
+            probs = cp.asarray(probabilities, dtype=cp.float64)
+            probs = probs / cp.sum(probs)
+
+        if replace:
+            if probs is None:
+                indices = gen.integers(0, n, size=target_size)
+            else:
+                cdf = cp.cumsum(probs)
+                draws = gen.random(size=target_size)
+                indices = cp.searchsorted(cdf, draws)
+        else:
+            flat_k = (
+                1
+                if target_size == ()
+                else int(cp.prod(cp.asarray(target_size)))
+            )
+            if probs is None:
+                indices = gen.permutation(n)[:flat_k]
+            else:
+                gumbels = -cp.log(-cp.log(gen.random(n)))
+                keys = cp.log(probs) + gumbels
+                indices = cp.argpartition(keys, -flat_k)[-flat_k:]
+            if target_size != ():
+                indices = cp.reshape(indices, target_size)
+            else:
+                indices = indices[0]
+
+        if values is None:
+            result = indices
+        else:
+            result = values[indices]
+        return self._result_and_next_state(gen, result)
+
+
+__all__ = ["CuPyBackend", "CuPyFunctionalBackend"]

orng/backends/jax.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import secrets
+from typing import Any
+
+from .._utils import SizeLike, normalize_shape
+
+
+class JAXBackend:
+    def __init__(self, *, seed: int | None, key: Any | None) -> None:
+        self._impl = JAXFunctionalBackend()
+        self._state = self._impl.init_state(seed=seed, generator=key)
+
+    def random(self, *, size: SizeLike, dtype: Any | None) -> Any:
+        result, self._state = self._impl.random(
+            self._state,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def uniform(
+        self,
+        *,
+        low: Any,
+        high: Any,
+        size: SizeLike,
+        dtype: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.uniform(
+            self._state,
+            low=low,
+            high=high,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def normal(
+        self,
+        *,
+        loc: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.normal(
+            self._state,
+            loc=loc,
+            scale=scale,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def gamma(
+        self,
+        *,
+        shape: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.gamma(
+            self._state,
+            shape=shape,
+            scale=scale,
+            size=size,
+            dtype=dtype,
+        )
+        return result
+
+    def choice(
+        self,
+        population: int | Any,
+        *,
+        size: SizeLike,
+        replace: bool,
+        probabilities: Any | None,
+    ) -> Any:
+        result, self._state = self._impl.choice(
+            self._state,
+            population,
+            size=size,
+            replace=replace,
+            probabilities=probabilities,
+        )
+        return result
+
+
+class JAXFunctionalBackend:
+    def __init__(self) -> None:
+        try:
+            import jax
+            import jax.numpy as jnp
+        except ImportError as exc:  # pragma: no cover - optional dependency
+            raise ImportError(
+                "JAX backend requires the 'jax' package to be installed. "
+                "Install it with `pip install orng[jax]`."
+            ) from exc
+
+        self._jax = jax
+        self._jnp = jnp
+
+    def init_state(self, *, seed: int | None, generator: Any | None) -> Any:
+        if generator is not None:
+            return generator
+        if seed is None:
+            seed = secrets.randbits(32)
+        return self._jax.random.key(seed)
+
+    def random(
+        self,
+        state: Any,
+        *,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        key, next_state = self._jax.random.split(state)
+        shape = normalize_shape(size)
+        sample_dtype = dtype if dtype is not None else self._jnp.float32
+        low = self._jnp.array(0.0, dtype=sample_dtype)
+        high = self._jnp.array(1.0, dtype=sample_dtype)
+        if shape:
+            result = self._jax.random.uniform(
+                key,
+                shape=shape,
+                minval=low,
+                maxval=high,
+                dtype=sample_dtype,
+            )
+        else:
+            result = self._jax.random.uniform(
+                key,
+                shape=(1,),
+                minval=low,
+                maxval=high,
+                dtype=sample_dtype,
+            )[0]
+        return result, next_state
+
+    def uniform(
+        self,
+        state: Any,
+        *,
+        low: Any,
+        high: Any,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        key, next_state = self._jax.random.split(state)
+        shape = normalize_shape(size)
+        sample_dtype = dtype if dtype is not None else self._jnp.float32
+        low_arr = self._jnp.asarray(low, dtype=sample_dtype)
+        high_arr = self._jnp.asarray(high, dtype=sample_dtype)
+        if shape:
+            result = self._jax.random.uniform(
+                key,
+                shape=shape,
+                minval=low_arr,
+                maxval=high_arr,
+                dtype=sample_dtype,
+            )
+        else:
+            result = self._jax.random.uniform(
+                key,
+                shape=(1,),
+                minval=low_arr,
+                maxval=high_arr,
+                dtype=sample_dtype,
+            )[0]
+        return result, next_state
+
+    def normal(
+        self,
+        state: Any,
+        *,
+        loc: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        key, next_state = self._jax.random.split(state)
+        shape = normalize_shape(size)
+        sample_dtype = dtype if dtype is not None else self._jnp.float32
+        if shape:
+            standard = self._jax.random.normal(
+                key,
+                shape=shape,
+                dtype=sample_dtype,
+            )
+        else:
+            standard = self._jax.random.normal(
+                key,
+                shape=(1,),
+                dtype=sample_dtype,
+            )[0]
+        return standard * scale + loc, next_state
+
+    def gamma(
+        self,
+        state: Any,
+        *,
+        shape: Any,
+        scale: Any,
+        size: SizeLike,
+        dtype: Any | None = None,
+    ) -> tuple[Any, Any]:
+        key, next_state = self._jax.random.split(state)
+        sample_shape = normalize_shape(size)
+        sample_dtype = dtype if dtype is not None else self._jnp.float32
+        concentration = self._jnp.asarray(shape, dtype=sample_dtype)
+        scale_arr = self._jnp.asarray(scale, dtype=sample_dtype)
+        if sample_shape:
+            draw_shape = sample_shape + self._jnp.shape(concentration)
+        else:
+            draw_shape = self._jnp.shape(concentration)
+        gamma_samples = self._jax.random.gamma(
+            key,
+            concentration,
+            shape=draw_shape,
+            dtype=sample_dtype,
+        )
+        scaled = gamma_samples * scale_arr
+        return scaled, next_state
+
+    def choice(
+        self,
+        state: Any,
+        population: int | Any,
+        *,
+        size: SizeLike,
+        replace: bool,
+        probabilities: Any | None,
+    ) -> tuple[Any, Any]:
+        key, next_state = self._jax.random.split(state)
+        shape = normalize_shape(size)
+        jax_shape = shape if shape else None
+        if isinstance(population, int):
+            domain = population
+        else:
+            domain = self._jnp.asarray(population)
+        probs = (
+            None if probabilities is None else self._jnp.asarray(probabilities)
+        )
+        result = self._jax.random.choice(
+            key,
+            domain,
+            shape=jax_shape,
+            replace=replace,
+            p=probs,
+        )
+        return result, next_state
+
+
+__all__ = ["JAXBackend", "JAXFunctionalBackend"]