flopscope 0.2.0__py3-none-any.whl

flopscope/numpy/random/_cost_formulas.py

@@ -0,0 +1,115 @@
+"""Named cost-formula vocabulary for fnp.random method-level entries.
+
+Each formula resolves to a callable ``(args, kwargs, result) -> int`` that
+computes the FLOP cost from the call arguments and the numpy result.
+The registry's ``cost_formula`` field names which formula a method uses.
+"""
+
+from __future__ import annotations
+
+import builtins as _builtins
+from collections.abc import Callable
+from typing import Any
+
+import numpy as _np
+
+from flopscope._flops import sort_cost as _sort_cost
+
+
+def _numel_output(args: tuple[Any, ...], kwargs: dict[str, Any], result: Any) -> int:
+    if isinstance(result, _np.ndarray):
+        return _builtins.max(int(result.size), 1)
+    return 1
+
+
+def _numel_input(args: tuple[Any, ...], kwargs: dict[str, Any], result: Any) -> int:
+    a = args[0] if args else kwargs.get("x")
+    if a is None:
+        return 1
+    if isinstance(a, _np.ndarray):
+        return _builtins.max(int(a.size), 1)
+    if hasattr(a, "__len__"):
+        return _builtins.max(len(a), 1)
+    return 1
+
+
+def _length(args: tuple[Any, ...], kwargs: dict[str, Any], result: Any) -> int:
+    if args:
+        n = int(args[0])
+    elif "length" in kwargs:
+        n = int(kwargs["length"])
+    else:
+        n = 1
+    return _builtins.max(n, 1)
+
+
+def _sort_cost_formula(
+    args: tuple[Any, ...], kwargs: dict[str, Any], result: Any
+) -> int:
+    a = args[0] if args else kwargs.get("a")
+    if a is None:
+        return _sort_cost(1)
+    if isinstance(a, (int, _np.integer)):
+        n = int(a)
+    elif isinstance(a, _np.ndarray):
+        n = int(a.shape[0]) if a.ndim > 0 else int(a)
+    elif hasattr(a, "__len__"):
+        n = len(a)
+    else:
+        n = 1
+    return _sort_cost(n)
+
+
+def _shape_axis(args: tuple[Any, ...], kwargs: dict[str, Any], result: Any) -> int:
+    """Cost = shape along the axis being permuted (defaults to axis=0).
+
+    Used by shuffle/permutation: the algorithm is O(shape[axis]) RNG draws
+    regardless of how wide each slice is. For integer input (the
+    ``permutation(int_n)`` case), cost = ``int(n)``. For ``axis=None`` —
+    which numpy interprets as "flatten then operate" — cost = numel.
+    """
+    a = args[0] if args else kwargs.get("x")
+    if a is None:
+        return 1
+    if isinstance(a, (int, _np.integer)):
+        return _builtins.max(int(a), 1)
+
+    axis = args[1] if len(args) >= 2 else kwargs.get("axis", 0)
+    if axis is None:
+        if isinstance(a, _np.ndarray):
+            return _builtins.max(int(a.size), 1)
+        if hasattr(a, "__len__"):
+            return _builtins.max(len(a), 1)
+        return 1
+
+    if isinstance(a, _np.ndarray):
+        if a.ndim == 0:
+            # 0-D scalar array; numpy choice/permutation treats as int(a)
+            return _builtins.max(int(a), 1)
+        return _builtins.max(int(a.shape[int(axis)]), 1)
+    if hasattr(a, "__len__"):
+        return _builtins.max(len(a), 1)
+    return 1
+
+
+def _choice_cost(args: tuple[Any, ...], kwargs: dict[str, Any], result: Any) -> int:
+    # Generator.choice:    choice(a, size=None, replace=True, p=None, axis=0, shuffle=True)
+    # RandomState.choice:  choice(a, size=None, replace=True, p=None)
+    # `replace` is the 3rd positional or the `replace` kwarg.
+    if len(args) >= 3:
+        replace = bool(args[2])
+    else:
+        replace = bool(kwargs.get("replace", True))
+    if replace:
+        return _numel_output(args, kwargs, result)
+    return _sort_cost_formula(args, kwargs, result)
+
+
+COST_FORMULAS: dict[str, Callable[[tuple[Any, ...], dict[str, Any], Any], int]] = {
+    "numel(output)": _numel_output,
+    "numel(input)": _numel_input,
+    "shape[axis]": _shape_axis,
+    "length": _length,
+    "sort_cost(n)": _sort_cost_formula,
+    "choice_cost": _choice_cost,
+}

flopscope/numpy/random/_counted_classes.py

@@ -0,0 +1,241 @@
+"""Counted subclasses of numpy.random.Generator and numpy.random.RandomState.
+
+The classes use ``__getattribute__`` as a gate that consults registry-derived
+``_COUNTED`` and ``_FREE`` sets. Sampler methods are added at module init by
+``_build_counted_class``, which iterates the registry and emits wrappers via
+``_make_counted_method``.
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+from typing import Any, ClassVar
+
+import numpy as _np
+
+from flopscope._budget import _counted_wrapper
+from flopscope._ndarray import _asflopscope
+from flopscope._registry import REGISTRY
+from flopscope._validation import require_budget
+from flopscope.errors import UnsupportedFunctionError
+from flopscope.numpy.random._cost_formulas import COST_FORMULAS
+
+
+@_counted_wrapper
+def _make_counted_method(
+    op_name: str, formula_name: str, base_cls: type, plain_factory: Callable[..., Any]
+) -> Callable[..., Any]:
+    """Build a counted-method wrapper for the parent class's named method.
+
+    The wrapper calls the parent method on a *plain* base-class instance that
+    shares the same RNG state (so internal sibling-method calls do not go
+    through the counted gate and double-count).  After numpy's C path returns,
+    cost is computed via the named formula, deducted from the active
+    BudgetContext, and the result is wrapped as FlopscopeArray if it is an
+    ndarray.
+    """
+    method_name = op_name.split(".")[-1]
+    base_method = getattr(base_cls, method_name)
+    formula = COST_FORMULAS[formula_name]
+
+    @functools.wraps(base_method)
+    @_counted_wrapper
+    def wrapped(self, *args: Any, **kwargs: Any) -> Any:
+        budget = require_budget()
+        # Dispatch through a plain base-class instance to avoid the counted
+        # __getattribute__ gate intercepting any internal sibling calls
+        # (e.g. Generator.choice → Generator.integers).
+        plain = plain_factory(self)
+        result = base_method(plain, *args, **kwargs)
+        cost = formula(args, kwargs, result)
+        with budget.deduct(op_name, flop_cost=cost, subscripts=None, shapes=()):
+            pass  # numpy already executed; deduct is post-hoc for output-dependent cost
+        if isinstance(result, _np.ndarray):
+            return _asflopscope(result)
+        return result
+
+    return wrapped
+
+
+def _build_counted_class(
+    base_cls: type,
+    op_prefix: str,
+    target_cls: type,
+    plain_factory: Callable[..., Any],
+) -> None:
+    """Populate target_cls._COUNTED, _FREE, and counted method bindings.
+
+    Reads REGISTRY entries whose op_name starts with op_prefix. counted_random_method
+    entries become wrapped class methods; free_random_method entries are added to
+    _FREE so the __getattribute__ gate lets them through unwrapped.
+
+    plain_factory is called with ``self`` (the counted subclass instance) to
+    produce a plain base-class instance sharing the same RNG state, so that
+    internal sibling calls inside numpy's C code do not go through the gate.
+    """
+    counted: set[str] = set()
+    free: set[str] = set()
+    for op_name, entry in REGISTRY.items():
+        if not op_name.startswith(op_prefix):
+            continue
+        short = op_name[len(op_prefix) :]
+        category = entry.get("category")
+        if category == "counted_random_method":
+            counted.add(short)
+            setattr(
+                target_cls,
+                short,
+                _make_counted_method(
+                    op_name, entry["cost_formula"], base_cls, plain_factory
+                ),
+            )
+        elif category == "free_random_method":
+            free.add(short)
+    target_cls._COUNTED = frozenset(counted)
+    target_cls._FREE = frozenset(free)
+
+
+def _reconstruct_counted_generator(bit_generator: Any) -> _CountedGenerator:
+    """Pickle helper — reconstruct a _CountedGenerator from its bit_generator."""
+    return _CountedGenerator(bit_generator)
+
+
+def _reconstruct_counted_random_state(state: tuple) -> _CountedRandomState:
+    """Pickle helper — reconstruct a _CountedRandomState from its full state tuple."""
+    # Construct with default seed; the state we just received will overwrite.
+    rs = _CountedRandomState(seed=None)
+    rs.set_state(state)
+    return rs
+
+
+class _CountedGenerator(_np.random.Generator):
+    """numpy ``Generator`` subclass with FLOP-counted sampler methods.
+
+    Each sampler method (``standard_normal``, ``normal``, ``uniform``,
+    ``integers``, ``choice``, ``shuffle``, ``permutation``, ``bytes``, ...)
+    deducts FLOPs from the active ``BudgetContext`` and returns
+    ``FlopscopeArray``. Free attribute access (``bit_generator``, ``spawn``)
+    is allowed; anything else raises ``UnsupportedFunctionError``.
+
+    Construct via :func:`flopscope.numpy.random.default_rng` (canonical) or by
+    passing a ``BitGenerator`` directly: ``Generator(np.random.PCG64(42))``.
+
+    Examples
+    --------
+    >>> import flopscope.numpy as fnp
+    >>> from flopscope import BudgetContext
+    >>> rng = fnp.random.default_rng(42)
+    >>> with BudgetContext(flop_budget=10**6):
+    ...     x = rng.standard_normal((10,))
+    >>> type(x).__name__
+    'FlopscopeArray'
+
+    Pickle / ``copy`` round-trips preserve counting:
+
+    >>> import pickle
+    >>> revived = pickle.loads(pickle.dumps(rng))
+    >>> isinstance(revived, type(rng))
+    True
+
+    See Also
+    --------
+    fnp.random.default_rng : canonical constructor.
+    """
+
+    _COUNTED: ClassVar[frozenset[str]] = frozenset()
+    _FREE: ClassVar[frozenset[str]] = frozenset()
+
+    def __getattribute__(self, name: str) -> Any:
+        if name.startswith("_"):
+            return super().__getattribute__(name)
+        cls = type(self)
+        if name in cls._FREE or name in cls._COUNTED:
+            return super().__getattribute__(name)
+        raise UnsupportedFunctionError(
+            f"flopscope does not count Generator.{name}. "
+            f"This is either a new numpy method or one not yet wrapped. "
+            f"See https://github.com/AIcrowd/flopscope/issues/18"
+        )
+
+    # Wrap spawned children so their methods stay counted.
+    # Annotated as list[Generator] to match parent's invariant return type;
+    # runtime objects are _CountedGenerator (subclass of Generator).
+    def spawn(self, n_children: int) -> list[_np.random.Generator]:
+        children = _np.random.Generator.spawn(self, n_children)
+        return [_CountedGenerator(c.bit_generator) for c in children]
+
+    # Round-trip pickling back to _CountedGenerator (numpy's default would
+    # reconstruct as np.random.Generator and lose the counting wrappers).
+    def __reduce__(self) -> tuple:
+        return (_reconstruct_counted_generator, (self.bit_generator,))
+
+
+class _CountedRandomState(_np.random.RandomState):
+    """numpy legacy ``RandomState`` subclass with FLOP-counted sampler methods.
+
+    Mirrors :class:`_CountedGenerator` for the legacy API: each sampler
+    (``randn``, ``normal``, ``uniform``, ``randint``, ``choice``,
+    ``shuffle``, ``permutation``, ``bytes``, ...) deducts FLOPs from the
+    active ``BudgetContext`` and returns ``FlopscopeArray``. Free methods
+    (``seed``, ``get_state``, ``set_state``) pass through; everything else
+    raises ``UnsupportedFunctionError``.
+
+    Modern code should prefer :func:`flopscope.numpy.random.default_rng`;
+    use this only when porting code that relies on the legacy API.
+
+    Examples
+    --------
+    >>> import flopscope.numpy as fnp
+    >>> from flopscope import BudgetContext
+    >>> rs = fnp.random.RandomState(42)
+    >>> with BudgetContext(flop_budget=10**6):
+    ...     z = rs.randn(10)
+    >>> type(z).__name__
+    'FlopscopeArray'
+
+    See Also
+    --------
+    fnp.random.default_rng : canonical (modern) RNG constructor.
+    """
+
+    _COUNTED: ClassVar[frozenset[str]] = frozenset()
+    _FREE: ClassVar[frozenset[str]] = frozenset()
+
+    def __getattribute__(self, name: str) -> Any:
+        if name.startswith("_"):
+            return super().__getattribute__(name)
+        cls = type(self)
+        if name in cls._FREE or name in cls._COUNTED:
+            return super().__getattribute__(name)
+        raise UnsupportedFunctionError(
+            f"flopscope does not count RandomState.{name}. "
+            f"This is either a new numpy method or one not yet wrapped. "
+            f"See https://github.com/AIcrowd/flopscope/issues/18"
+        )
+
+    # Round-trip pickling back to _CountedRandomState (numpy's default
+    # __reduce__ hardcodes __randomstate_ctor which returns plain RandomState,
+    # silently losing FLOP counting).
+    def __reduce__(self) -> tuple:
+        # get_state() is in _FREE so the gate lets it through.
+        # set_state() is also in _FREE, used by the helper on the other side.
+        return (_reconstruct_counted_random_state, (self.get_state(),))
+
+
+# Wire counted methods at import time.
+# plain_factory creates a plain base-class instance sharing the same RNG state;
+# this prevents numpy's internal sibling-method calls from going through the
+# counted __getattribute__ gate and double-counting FLOPs.
+_build_counted_class(
+    _np.random.Generator,
+    "random.Generator.",
+    _CountedGenerator,
+    plain_factory=lambda self: _np.random.Generator(self._bit_generator),
+)
+_build_counted_class(
+    _np.random.RandomState,
+    "random.RandomState.",
+    _CountedRandomState,
+    plain_factory=lambda self: _np.random.RandomState(self._bit_generator),
+)

flopscope/numpy/testing/__init__.py

@@ -0,0 +1,13 @@
+"""Test utilities re-exported from ``numpy.testing``.
+
+These are free operations — they perform assertions and comparisons
+for use in tests and profiling code, and do not count against any
+FLOP budget.
+"""
+
+from numpy.testing import assert_allclose, assert_array_equal  # noqa: F401
+
+__all__ = [
+    "assert_allclose",
+    "assert_array_equal",
+]

flopscope/numpy/typing/__init__.py

@@ -0,0 +1,30 @@
+"""Type hints re-exported from numpy.typing.
+
+Usage::
+
+    from flopscope.numpy.typing import NDArray, ArrayLike
+
+    def predict(x: NDArray) -> NDArray:
+        ...
+
+All names are plain aliases to numpy.typing. Because FlopscopeArray
+is a subclass of numpy.ndarray, annotations like `NDArray[float32]`
+accept flopscope arrays without any runtime overhead.
+"""
+
+from __future__ import annotations
+
+# Re-export everything numpy.typing publicly exposes
+from numpy import typing as _np_typing
+from numpy.typing import (  # noqa: F401
+    ArrayLike,
+    DTypeLike,
+    NBitBase,
+    NDArray,
+)
+
+__all__ = [name for name in dir(_np_typing) if not name.startswith("_")]
+
+for _name in __all__:
+    if _name not in globals():
+        globals()[_name] = getattr(_np_typing, _name)

flopscope/stats/__init__.py

@@ -0,0 +1,84 @@
+"""Continuous probability distributions with analytic FLOP accounting.
+
+``flopscope.stats`` provides a focused subset of ``scipy.stats`` continuous
+distributions. Each exported distribution object exposes ``pdf``, ``cdf``,
+and ``ppf`` methods with SciPy-compatible signatures while charging a flat
+FLOP cost per output element to the active budget.
+
+Available distributions
+-----------------------
+norm
+    Normal (Gaussian) distribution.
+uniform
+    Continuous uniform distribution.
+expon
+    Exponential distribution.
+cauchy
+    Cauchy (Lorentz) distribution.
+logistic
+    Logistic distribution.
+laplace
+    Laplace (double-exponential) distribution.
+lognorm
+    Log-normal distribution.
+truncnorm
+    Truncated normal distribution.
+
+Notes
+-----
+All distributions use SciPy's ``loc``/``scale`` parameterization. Shape
+parameters, when present, precede ``loc`` and ``scale`` exactly as they do in
+``scipy.stats``. Each public method requires an active
+``flopscope.BudgetContext`` and deducts the documented flat FLOP charge before
+returning a ``FlopscopeArray`` result.
+
+    import flopscope as flops
+
+    flops.stats.norm.pdf(0.0)                    # standard normal PDF at 0
+    flops.stats.norm.cdf(1.96, loc=0, scale=1)   # ≈ 0.975
+    flops.stats.expon.ppf(0.5, scale=2.0)        # median of Exp(rate=0.5)
+
+FLOP costs
+----------
+Each method deducts a **flat FLOP cost per element** from the active budget.
+Costs are documented in each method's docstring. Internal sub-operations
+(exp, log, etc.) are computed via raw NumPy and do **not** incur additional
+FLOP charges.
+
+Compatibility
+-------------
+Outputs are verified against ``scipy.stats`` to within 1e-12 relative
+tolerance across the full input domain. See ``tests/test_stats_*.py``.
+
+Examples
+--------
+>>> import flopscope as flops
+>>> with flops.BudgetContext(flop_budget=32) as budget:
+...     probs = flops.stats.norm.cdf([0.0, 1.96])
+...     summary = budget.summary()
+"""
+
+from flopscope._registry import make_module_getattr as _make_module_getattr
+from flopscope.stats._cauchy import cauchy  # noqa: F401
+from flopscope.stats._expon import expon  # noqa: F401
+from flopscope.stats._laplace import laplace  # noqa: F401
+from flopscope.stats._logistic import logistic  # noqa: F401
+from flopscope.stats._lognorm import lognorm  # noqa: F401
+from flopscope.stats._norm import norm  # noqa: F401
+from flopscope.stats._truncnorm import truncnorm  # noqa: F401
+from flopscope.stats._uniform import uniform  # noqa: F401
+
+__all__ = [
+    "cauchy",
+    "expon",
+    "laplace",
+    "logistic",
+    "lognorm",
+    "norm",
+    "truncnorm",
+    "uniform",
+]
+
+__getattr__ = _make_module_getattr(
+    module_prefix="stats.", module_label="flopscope.stats"
+)

flopscope/stats/_base.py

@@ -0,0 +1,77 @@
+"""Base support for SciPy-compatible continuous distributions."""
+
+from __future__ import annotations
+
+import numpy as _np
+
+from flopscope._budget import _counted_wrapper
+from flopscope._ndarray import _asflopscope
+from flopscope._validation import require_budget
+
+
+class ContinuousDistribution:
+    """Base class for FLOP-counted continuous distributions.
+
+    Parameters
+    ----------
+    name : str
+        Distribution name used to construct operation labels such as
+        ``"stats.norm.pdf"`` in the budget log.
+
+    Notes
+    -----
+    Subclasses implement ``_compute_pdf``, ``_compute_cdf``, and
+    ``_compute_ppf`` as pure NumPy kernels. Public ``pdf``, ``cdf``, and
+    ``ppf`` methods should delegate through :meth:`_deduct_and_call` so that
+    budget deduction and ``FlopscopeArray`` wrapping stay consistent across
+    the stats surface.
+    """
+
+    def __init__(self, name: str):
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @_counted_wrapper
+    def _deduct_and_call(self, method: str, cost_per_elem: int, x, *args, **kwargs):
+        """Deduct FLOPs then call the pure-numpy implementation.
+
+        Parameters
+        ----------
+        method : str
+            Method name for budget logging, e.g. ``"pdf"``.
+        cost_per_elem : int
+            Flat FLOP cost per output element.
+        x : array_like
+            Primary input array (determines output size).
+        *args, **kwargs
+            Forwarded to ``_compute_{method}``.
+
+        Returns
+        -------
+        FlopscopeArray
+            Result returned by the matching ``_compute_{method}``
+            implementation after budget deduction.
+
+        Notes
+        -----
+        The deducted FLOP charge is ``cost_per_elem * max(numel(x), 1)``.
+        """
+        budget = require_budget()
+        x = _np.asarray(x, dtype=_np.float64)
+        n = max(x.size, 1)
+        op_name = f"stats.{self._name}.{method}"
+        compute_fn = getattr(self, f"_compute_{method}")
+        with budget.deduct(
+            op_name,
+            flop_cost=cost_per_elem * n,
+            subscripts=None,
+            shapes=(x.shape,),
+        ):
+            result = compute_fn(x, *args, **kwargs)
+        return _asflopscope(result)
+
+    def __repr__(self) -> str:
+        return f"<flopscope.stats.{self._name}>"