flopscope 0.2.0__py3-none-any.whl

flopscope/_config.py

@@ -0,0 +1,118 @@
+"""Global configuration for flopscope."""
+
+from __future__ import annotations
+
+_SETTINGS: dict[str, object] = {
+    "check_nan_inf": False,
+    "dimino_budget": 50_000,
+    "einsum_path_cache_size": 4096,
+    "partition_budget": 100_000,
+    "symmetry_warnings": True,
+}
+
+# Validators for settings that require range/type checks.
+# Each validator receives the proposed value and raises ValueError/TypeError
+# if the value is invalid.
+_VALIDATORS: dict[str, object] = {
+    "dimino_budget": lambda v: _require_non_negative_int("dimino_budget", v),
+    "partition_budget": lambda v: _require_non_negative_int("partition_budget", v),
+}
+
+
+def _require_non_negative_int(name: str, value: object) -> None:
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise TypeError(
+            f"Setting {name!r} requires a non-negative int; got {type(value).__name__!r}"
+        )
+    if value < 0:
+        raise ValueError(f"Setting {name!r} requires a non-negative int; got {value!r}")
+
+
+def configure(**kwargs: object) -> None:
+    """Update flopscope global settings.
+
+    Parameters
+    ----------
+    check_nan_inf : bool
+        If ``True``, scan every counted op's output for NaN/Inf values and
+        emit a :class:`~flopscope.errors.FlopscopeWarning` if any are found.
+        The scan is two full O(n) sweeps over the result and is attributed
+        to ``flopscope_overhead_time``, so it is off by default for
+        production scoring.  Opt in when debugging an estimator that
+        produces NaN/Inf to identify the introducing op.  Default ``False``.
+    dimino_budget : int
+        Maximum number of group elements during whole-expression G_pt closure.
+        Pathological declared-symmetry cases that exceed this budget fall back
+        to dense cost with a CostFallbackWarning.  Default ``50_000`` — accepts
+        groups up to ``|S_8| = 40_320`` with a small margin; rejects ``|S_9|``
+        and larger.  Tune via ``flops.configure(dimino_budget=...)``;
+        ``benchmarks/_perm_group_calibration.py`` recommends a value for the
+        host machine.
+    einsum_path_cache_size : int
+        Maximum number of entries in the einsum path cache.
+        Changing this rebuilds the cache (old entries are discarded).
+        Default ``4096``.
+    partition_budget : int
+        Maximum number of typed partitions a single component may have before
+        the partitionCount regime refuses.  Components exceeding this budget
+        fall back to the dense cost with a CostFallbackWarning.  Default
+        ``100_000`` (covers Bell(9)=21_147; Bell(10)=115_975 forces fallback).
+        Set to ``0`` to force fallback for any non-trivial component.
+    symmetry_warnings : bool
+        If ``False``, suppress :class:`~flopscope.errors.SymmetryLossWarning`
+        warnings.  Default ``True``.
+
+    Returns
+    -------
+    None
+        Updates the in-process global configuration immediately.
+
+    Examples
+    --------
+    >>> import flopscope as flops
+    >>> flops.configure(einsum_path_cache_size=8192)
+    >>> flops.configure(symmetry_warnings=False)
+    >>> flops.configure(check_nan_inf=True)
+    >>> flops.configure(partition_budget=50_000)
+    >>> flops.configure(dimino_budget=1_000_000)
+    """
+    for key, value in kwargs.items():
+        if key not in _SETTINGS:
+            raise ValueError(f"Unknown setting: {key!r}")
+        validator = _VALIDATORS.get(key)
+        if validator is not None:
+            validator(value)  # type: ignore[call-arg]
+        _SETTINGS[key] = value
+
+    if "einsum_path_cache_size" in kwargs:
+        from flopscope._einsum import _rebuild_einsum_cache
+
+        _rebuild_einsum_cache()
+
+
+def get_setting(key: str) -> object:
+    """Return the current value of a global setting."""
+    return _SETTINGS[key]
+
+
+def set_setting(key: str, value: object) -> None:
+    """Set a single global setting by name.
+
+    Equivalent to ``configure(**{key: value})``.  Provided for call-sites
+    that hold the setting name in a variable.
+
+    Parameters
+    ----------
+    key : str
+        The setting name.
+    value : object
+        The new value.  Subject to the same validation as :func:`configure`.
+
+    Raises
+    ------
+    ValueError
+        If *key* is unknown or the value fails range validation.
+    TypeError
+        If the value fails type validation.
+    """
+    configure(**{key: value})

flopscope/_counting_ops.py

@@ -0,0 +1,451 @@
+"""Counted wrappers for trace, histogram, and generation operations.
+
+These are operations that look "free" but involve genuine computation.
+Each function charges a FLOP cost to the active budget.
+"""
+
+from __future__ import annotations
+
+import builtins as _builtins
+import inspect as _inspect
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as _np
+from numpy.typing import ArrayLike, DTypeLike
+
+from flopscope._budget import _call_numpy, _counted_wrapper
+from flopscope._docstrings import attach_docstring
+from flopscope._flops import _ceil_log2
+from flopscope._ndarray import FlopscopeArray, _to_base_ndarray, _to_base_ndarray_tree
+from flopscope._validation import require_budget
+
+# ---------------------------------------------------------------------------
+# Reductions disguised as free
+# ---------------------------------------------------------------------------
+
+
+@_counted_wrapper
+def trace(
+    a: ArrayLike,
+    offset: int = 0,
+    axis1: int = 0,
+    axis2: int = 1,
+    dtype: DTypeLike | None = None,
+    out: FlopscopeArray | None = None,
+) -> FlopscopeArray:
+    budget = require_budget()
+    a = _np.asarray(a)
+    cost = _builtins.max(_builtins.min(a.shape[axis1], a.shape[axis2]), 1)
+    out_stripped = _to_base_ndarray(out) if out is not None else None
+    with budget.deduct("trace", flop_cost=cost, subscripts=None, shapes=(a.shape,)):
+        result = _call_numpy(
+            _np.trace,
+            _to_base_ndarray(a),
+            offset=offset,
+            axis1=axis1,
+            axis2=axis2,
+            dtype=dtype,
+            out=out_stripped,
+        )
+    return out if out is not None else result  # type: ignore[return-value]
+
+
+attach_docstring(
+    trace, _np.trace, "counted_custom", "min(a.shape[axis1], a.shape[axis2]) FLOPs"
+)
+
+
+@_counted_wrapper
+def allclose(a: ArrayLike, b: ArrayLike, **kwargs: Any) -> bool:
+    budget = require_budget()
+    a = _np.asarray(a)
+    b = _np.asarray(b)
+    out_shape = _np.broadcast_shapes(a.shape, b.shape)
+    numel = 1
+    for d in out_shape:
+        numel *= d
+    cost = _builtins.max(numel, 1)
+    with budget.deduct(
+        "allclose", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape)
+    ):
+        result = _call_numpy(_np.allclose, a, b, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(allclose, _np.allclose, "counted_custom", "numel(a) FLOPs")
+allclose.__signature__ = _inspect.signature(_np.allclose)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def array_equal(a: ArrayLike, b: ArrayLike, **kwargs: Any) -> bool:
+    budget = require_budget()
+    a = _np.asarray(a)
+    b = _np.asarray(b)
+    # array_equal does not broadcast; returns False on shape mismatch
+    cost = _builtins.max(a.size, b.size, 1)
+    with budget.deduct(
+        "array_equal", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape)
+    ):
+        result = _call_numpy(_np.array_equal, a, b, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(array_equal, _np.array_equal, "counted_custom", "numel(a) FLOPs")
+array_equal.__signature__ = _inspect.signature(_np.array_equal)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def array_equiv(a: ArrayLike, b: ArrayLike) -> bool:
+    budget = require_budget()
+    a = _np.asarray(a)
+    b = _np.asarray(b)
+    # array_equiv broadcasts; returns False if shapes are incompatible
+    try:
+        out_shape = _np.broadcast_shapes(a.shape, b.shape)
+        numel = 1
+        for d in out_shape:
+            numel *= d
+        cost = _builtins.max(numel, 1)
+    except ValueError:
+        cost = _builtins.max(a.size, b.size, 1)
+    with budget.deduct(
+        "array_equiv", flop_cost=cost, subscripts=None, shapes=(a.shape, b.shape)
+    ):
+        result = _call_numpy(_np.array_equiv, a, b)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(array_equiv, _np.array_equiv, "counted_custom", "numel(a) FLOPs")
+array_equiv.__signature__ = _inspect.signature(_np.array_equiv)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+# ---------------------------------------------------------------------------
+# Histogram & counting
+# ---------------------------------------------------------------------------
+
+
+@_counted_wrapper
+def histogram(
+    a: ArrayLike,
+    bins: int | Sequence[int] | str = 10,
+    **kwargs: Any,
+) -> tuple[FlopscopeArray, FlopscopeArray]:
+    budget = require_budget()
+    a = _np.asarray(a)
+    n = a.size
+    if isinstance(bins, _builtins.int):
+        cost = _builtins.max(n * _ceil_log2(bins), 1)
+    elif isinstance(bins, _builtins.str):
+        cost = _builtins.max(n, 1)
+    else:
+        bins_arr = _np.asarray(bins)
+        cost = _builtins.max(n * _ceil_log2(_builtins.len(bins_arr)), 1)
+    with budget.deduct("histogram", flop_cost=cost, subscripts=None, shapes=(a.shape,)):
+        result = _call_numpy(_np.histogram, a, bins=bins, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(
+    histogram,
+    _np.histogram,
+    "counted_custom",
+    "n * ceil(log2(bins)) FLOPs when bins is int; n FLOPs otherwise",
+)
+histogram.__signature__ = _inspect.signature(_np.histogram)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def histogram2d(
+    x: ArrayLike,
+    y: ArrayLike,
+    bins: Any = 10,
+    **kwargs: Any,
+) -> tuple[FlopscopeArray, FlopscopeArray, FlopscopeArray]:
+    budget = require_budget()
+    x = _np.asarray(x)
+    y = _np.asarray(y)
+    n = x.size
+
+    # Determine bx, by
+    if isinstance(bins, _builtins.int):
+        cost = _builtins.max(n * (_ceil_log2(bins) + _ceil_log2(bins)), 1)
+    elif (
+        isinstance(bins, (_builtins.list, tuple))
+        and _builtins.len(bins) == 2
+        and isinstance(bins[0], _builtins.int)
+        and isinstance(bins[1], _builtins.int)
+    ):
+        bx, by = bins[0], bins[1]
+        cost = _builtins.max(n * (_ceil_log2(bx) + _ceil_log2(by)), 1)
+    elif isinstance(bins, (_builtins.list, tuple)) and _builtins.len(bins) == 2:
+        b0 = _np.asarray(bins[0])
+        b1 = _np.asarray(bins[1])
+        if b0.ndim >= 1 and b1.ndim >= 1:
+            cost = _builtins.max(
+                n * (_ceil_log2(_builtins.len(b0)) + _ceil_log2(_builtins.len(b1))), 1
+            )
+        else:
+            cost = _builtins.max(n, 1)
+    else:
+        cost = _builtins.max(n, 1)
+
+    with budget.deduct(
+        "histogram2d", flop_cost=cost, subscripts=None, shapes=(x.shape, y.shape)
+    ):
+        result = _call_numpy(_np.histogram2d, x, y, bins=bins, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(
+    histogram2d,
+    _np.histogram2d,
+    "counted_custom",
+    "n * (ceil(log2(bx)) + ceil(log2(by))) FLOPs when bins is int pair; n FLOPs otherwise",
+)
+histogram2d.__signature__ = _inspect.signature(_np.histogram2d)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def histogramdd(
+    sample: ArrayLike,
+    bins: Any = 10,
+    **kwargs: Any,
+) -> tuple[FlopscopeArray, list[FlopscopeArray]]:
+    budget = require_budget()
+    sample = _np.asarray(sample)
+    # sample shape: (n, d) or (n,) for 1-d
+    if sample.ndim == 1:
+        n = sample.shape[0]
+        d = 1
+    else:
+        n, d = sample.shape[0], sample.shape[1]
+
+    if isinstance(bins, _builtins.int):
+        cost = _builtins.max(n * d * _ceil_log2(bins), 1)
+    elif isinstance(bins, (_builtins.list, tuple)):
+        total_log = 0
+        for b in bins:
+            if isinstance(b, _builtins.int):
+                total_log += _ceil_log2(b)
+            else:
+                b_arr = _np.asarray(b)
+                if b_arr.ndim >= 1 and b_arr.size > 0:
+                    total_log += _ceil_log2(_builtins.len(b_arr))
+                else:
+                    total_log += 1
+        cost = _builtins.max(n * total_log, 1)
+    else:
+        cost = _builtins.max(n, 1)
+
+    with budget.deduct(
+        "histogramdd", flop_cost=cost, subscripts=None, shapes=(sample.shape,)
+    ):
+        result = _call_numpy(_np.histogramdd, sample, bins=bins, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(
+    histogramdd,
+    _np.histogramdd,
+    "counted_custom",
+    "n * d * ceil(log2(bins)) FLOPs when bins is int; n FLOPs otherwise",
+)
+histogramdd.__signature__ = _inspect.signature(_np.histogramdd)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def histogram_bin_edges(
+    a: ArrayLike,
+    bins: int | Sequence[int] | str = 10,
+    **kwargs: Any,
+) -> FlopscopeArray:
+    budget = require_budget()
+    a = _np.asarray(a)
+    cost = _builtins.max(a.size, 1)
+    with budget.deduct(
+        "histogram_bin_edges", flop_cost=cost, subscripts=None, shapes=(a.shape,)
+    ):
+        result = _call_numpy(_np.histogram_bin_edges, a, bins=bins, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(
+    histogram_bin_edges, _np.histogram_bin_edges, "counted_custom", "numel(a) FLOPs"
+)
+histogram_bin_edges.__signature__ = _inspect.signature(_np.histogram_bin_edges)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def bincount(x: ArrayLike, **kwargs: Any) -> FlopscopeArray:
+    budget = require_budget()
+    x = _np.asarray(x)
+    cost = _builtins.max(x.size, 1)
+    with budget.deduct("bincount", flop_cost=cost, subscripts=None, shapes=(x.shape,)):
+        result = _call_numpy(_np.bincount, x, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(bincount, _np.bincount, "counted_custom", "numel(x) FLOPs")
+try:
+    bincount.__signature__ = _inspect.signature(_np.bincount)  # pyright: ignore[reportFunctionMemberAccess]
+except (ValueError, TypeError):
+    pass
+
+
+# ---------------------------------------------------------------------------
+# Array generation with arithmetic
+# ---------------------------------------------------------------------------
+
+
+@_counted_wrapper
+def logspace(
+    start: Any,
+    stop: Any,
+    num: int = 50,
+    **kwargs: Any,
+) -> FlopscopeArray:
+    budget = require_budget()
+    cost = _builtins.max(num, 1)
+    with budget.deduct("logspace", flop_cost=cost, subscripts=None, shapes=((num,),)):
+        result = _call_numpy(_np.logspace, start, stop, num=num, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(logspace, _np.logspace, "counted_custom", "num FLOPs")
+logspace.__signature__ = _inspect.signature(_np.logspace)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def geomspace(
+    start: Any,
+    stop: Any,
+    num: int = 50,
+    **kwargs: Any,
+) -> FlopscopeArray:
+    budget = require_budget()
+    cost = _builtins.max(num, 1)
+    with budget.deduct("geomspace", flop_cost=cost, subscripts=None, shapes=((num,),)):
+        result = _call_numpy(_np.geomspace, start, stop, num=num, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(geomspace, _np.geomspace, "counted_custom", "num FLOPs")
+geomspace.__signature__ = _inspect.signature(_np.geomspace)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+@_counted_wrapper
+def vander(
+    x: ArrayLike,
+    N: int | None = None,
+    **kwargs: Any,
+) -> FlopscopeArray:
+    budget = require_budget()
+    x = _np.asarray(x)
+    n = _builtins.len(x)
+    if N is None:
+        N = n
+    cost = _builtins.max(n * (N - 1), 1)
+    with budget.deduct("vander", flop_cost=cost, subscripts=None, shapes=(x.shape,)):
+        result = _call_numpy(_np.vander, x, N=N, **kwargs)
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(vander, _np.vander, "counted_custom", "len(x) * (N-1) FLOPs")
+vander.__signature__ = _inspect.signature(_np.vander)  # pyright: ignore[reportFunctionMemberAccess]
+
+# ---------------------------------------------------------------------------
+# Apply & piecewise (formerly blacklisted)
+# ---------------------------------------------------------------------------
+
+
+@_counted_wrapper
+def apply_along_axis(
+    func1d: Any,
+    axis: int,
+    arr: ArrayLike,
+    *args: Any,
+    **kwargs: Any,
+) -> FlopscopeArray:
+    """Counted version of np.apply_along_axis. Cost: numel(output)."""
+    budget = require_budget()
+    if not isinstance(arr, _np.ndarray):
+        arr = _np.asarray(arr)
+    result = _np.apply_along_axis(func1d, axis, _to_base_ndarray(arr), *args, **kwargs)
+    cost = result.size if hasattr(result, "size") else 1
+    with budget.deduct(
+        "apply_along_axis", flop_cost=cost, subscripts=None, shapes=(arr.shape,)
+    ):
+        pass
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(
+    apply_along_axis,
+    _np.apply_along_axis,
+    "counted_custom",
+    "numel(output) FLOPs",
+)
+
+
+@_counted_wrapper
+def apply_over_axes(
+    func: Any,
+    a: ArrayLike,
+    axes: int | Sequence[int],
+) -> FlopscopeArray:
+    """Counted version of np.apply_over_axes. Cost: numel(output)."""
+    budget = require_budget()
+    if not isinstance(a, _np.ndarray):
+        a = _np.asarray(a)
+    result = _np.apply_over_axes(func, _to_base_ndarray(a), axes)
+    cost = result.size if hasattr(result, "size") else 1
+    with budget.deduct(
+        "apply_over_axes", flop_cost=cost, subscripts=None, shapes=(a.shape,)
+    ):
+        pass
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(
+    apply_over_axes,
+    _np.apply_over_axes,
+    "counted_custom",
+    "numel(output) FLOPs",
+)
+
+
+@_counted_wrapper
+def piecewise(
+    x: ArrayLike,
+    condlist: Any,
+    funclist: Any,
+    *args: Any,
+    **kw: Any,
+) -> FlopscopeArray:
+    """Counted version of np.piecewise. Cost: numel(input)."""
+    budget = require_budget()
+    if not isinstance(x, _np.ndarray):
+        x = _np.asarray(x)
+    result = _np.piecewise(
+        _to_base_ndarray(x), _to_base_ndarray_tree(condlist), funclist, *args, **kw
+    )
+    cost = x.size
+    with budget.deduct("piecewise", flop_cost=cost, subscripts=None, shapes=(x.shape,)):
+        pass
+    return result  # type: ignore[return-value]
+
+
+attach_docstring(
+    piecewise,
+    _np.piecewise,
+    "counted_custom",
+    "numel(input) FLOPs",
+)
+
+
+import sys as _sys  # noqa: E402
+
+from flopscope._ndarray import wrap_module_returns as _wrap_module_returns  # noqa: E402
+
+_wrap_module_returns(_sys.modules[__name__])