flopscope 0.2.0__py3-none-any.whl

flopscope/__init__.py

@@ -0,0 +1,273 @@
+"""Flopscope: FLOP-counting numpy primitives for the Mechanistic Estimation Challenge.
+
+The public API is structured JAX-style:
+
+- ``flopscope`` (this module) exports flopscope-specific primitives: budget
+  contexts, configuration, symmetric tensors, the unified
+  :class:`SymmetryGroup`, errors, and the :class:`FlopscopeArray` type.
+- ``flopscope.numpy`` exports the full counted numpy-shaped surface
+  (``einsum``, ``array``, ``linspace``, ``linalg``, ``fft``, ``random``, ...).
+  Attributes not explicitly implemented there fall back to raw ``numpy``.
+- ``flopscope.accounting`` exposes analytical cost helpers
+  (``einsum_cost``, ``pointwise_cost``, ``reduction_cost``, ...).
+- ``flopscope.stats`` hosts statistical-distribution primitives (not a
+  numpy submodule; closer in spirit to scipy.stats).
+
+Usage::
+
+    import flopscope as flops
+    import flopscope.numpy as fnp
+
+    flops.configure(symmetry_warnings=False)
+    with flops.BudgetContext(flop_budget=1_000_000) as budget:
+        W = fnp.array(weight_matrix)
+        h = fnp.einsum('ij,j->i', W, x)
+        h = fnp.maximum(h, 0)
+        print(budget.summary())
+"""
+
+import importlib as _importlib
+
+import numpy as _np
+
+from flopscope._registry import REGISTRY_META as _REGISTRY_META
+
+__version__ = f"0.2.0+np{_np.__version__}"
+__numpy_version__ = _np.__version__
+__numpy_pinned__ = _REGISTRY_META["numpy_version"]
+__numpy_supported__ = _REGISTRY_META.get("numpy_supported", ">=2.0.0,<2.5.0")
+
+from flopscope._version_check import check_numpy_version as _check_numpy_version
+
+_check_numpy_version(__numpy_supported__)
+
+# --- Symmetry-aware accumulation cost (einsum + reduction) ---
+from flopscope._accumulation import (  # noqa: F401,E402
+    AccumulationCost,
+    ComponentCost,
+    RegimeStep,
+    einsum_accumulation_cost,
+    reduction_accumulation_cost,
+)
+
+# --- Budget and diagnostics ---
+from flopscope._budget import (  # noqa: F401,E402
+    BudgetContext,
+    OpRecord,
+    budget,
+    budget_reset,
+    budget_summary_dict,
+    namespace,
+)
+from flopscope._config import configure  # noqa: F401,E402
+from flopscope._display import budget_live, budget_summary  # noqa: F401,E402
+
+# --- Array type (flopscope-specific) ---
+from flopscope._ndarray import FlopscopeArray  # noqa: F401,E402
+
+# --- Path optimization types ---
+from flopscope._opt_einsum import PathInfo, StepInfo  # noqa: F401,E402
+
+# --- Symmetry (post-#51 unified surface) ---
+from flopscope._perm_group import SymmetryGroup  # noqa: F401,E402
+
+# --- Symmetric tensor ---
+from flopscope._symmetric import (  # noqa: F401,E402
+    SymmetricTensor,
+    as_symmetric,
+    is_symmetric,
+    symmetrize,
+)
+
+# --- Errors ---
+from flopscope.errors import (  # noqa: F401,E402
+    BudgetExhaustedError,
+    FlopscopeError,
+    FlopscopeWarning,
+    NoBudgetContextError,
+    SymmetryError,
+    SymmetryLossWarning,
+    TimeExhaustedError,
+    UnsupportedFunctionError,
+)
+
+
+def einsum_clear_caches() -> None:
+    """Clear all flopscope einsum-related LRU caches.
+
+    Clears both the einsum path cache (consulted by ``fnp.einsum`` and
+    ``fnp.einsum_path``) and the einsum accumulation-cost cache (consulted
+    by ``fnp.einsum`` and ``flopscope.einsum_accumulation_cost``).
+
+    Useful when benchmarking cold-call latency. ``fnp.clear_einsum_cache``
+    still exists and clears only the path cache.
+    """
+    from flopscope._accumulation._cache import _accumulation_cache
+    from flopscope._einsum import _path_cache
+
+    _path_cache.cache_clear()
+    _accumulation_cache.cache_clear()
+
+
+def einsum_cache_info() -> dict:
+    """Return cache statistics for the einsum path + accumulation caches.
+
+    Returns
+    -------
+    dict
+        ``{"path": CacheInfo, "accumulation": CacheInfo}`` where each value
+        is a standard ``functools.lru_cache`` info tuple with ``hits``,
+        ``misses``, ``maxsize``, and ``currsize``.
+    """
+    from flopscope._accumulation._cache import _accumulation_cache
+    from flopscope._einsum import _path_cache
+
+    return {
+        "path": _path_cache.cache_info(),
+        "accumulation": _accumulation_cache.cache_info(),
+    }
+
+
+def reduction_clear_cache() -> None:
+    """Clear the reduction accumulation-cost cache.
+
+    Consulted by ``flopscope.reduction_accumulation_cost`` and the
+    ``np.ufunc.reduce`` family wrappers (``fnp.sum``, ``fnp.mean``,
+    ``fnp.median``, etc.).
+    """
+    from flopscope._accumulation._cache import _reduction_cache
+
+    _reduction_cache.cache_clear()
+
+
+def reduction_cache_info():
+    """Return cache statistics for the reduction accumulation-cost cache.
+
+    Returns a standard ``functools.lru_cache`` info tuple with ``hits``,
+    ``misses``, ``maxsize``, and ``currsize``.
+    """
+    from flopscope._accumulation._cache import _reduction_cache
+
+    return _reduction_cache.cache_info()
+
+
+def clear_cache() -> None:
+    """Clear all flopscope LRU caches (einsum + reduction).
+
+    Convenience aggregate over :func:`einsum_clear_caches` and
+    :func:`reduction_clear_cache`. Use the per-domain variants if you
+    only need to invalidate one cache.
+    """
+    einsum_clear_caches()
+    reduction_clear_cache()
+
+
+def tier2_reduction_cost(a, axis=None, *, dense_per_output_cost=None):
+    """Cost for selection-style reductions (median, percentile, quantile).
+
+    Returns ``num_output_orbits × dense_per_output_cost`` — one partition
+    pass per unique output cell. Under symmetry, orbit-shared output
+    cells share the partition pass; ``num_output_orbits`` is computed
+    from the input's declared symmetry.
+
+    Parameters
+    ----------
+    a : numpy.ndarray or SymmetricTensor
+        Input array. ``SymmetricTensor`` inputs contribute their declared
+        symmetry to the orbit count.
+    axis : int, tuple of int, or None
+        Axes being reduced. ``None`` (default) reduces all axes.
+    dense_per_output_cost : int, optional
+        The dense per-output cost. If ``None`` (the default), it is set
+        to the product of the reduced axes' lengths — i.e., one
+        partition pass per output cell, which is how ``fnp.median``,
+        ``fnp.percentile``, and ``fnp.quantile`` are charged. Pass an
+        explicit value to model a custom Tier-2 selection cost.
+
+    Returns
+    -------
+    int
+        Total flops charged.
+    """
+    import math as _math
+
+    import numpy as _np
+
+    from flopscope._pointwise import _tier2_reduction_cost
+
+    if not isinstance(a, _np.ndarray):
+        a = _np.asarray(a)
+
+    if dense_per_output_cost is None:
+        if axis is None:
+            dense_per_output_cost = _math.prod(a.shape) if a.shape else 1
+        else:
+            axes = (axis,) if isinstance(axis, int) else tuple(axis)
+            dense_per_output_cost = _math.prod(a.shape[i] for i in axes)
+
+    return _tier2_reduction_cost(a, axis, dense_per_output_cost)
+
+
+_LAZY_SUBMODULES = frozenset({"numpy", "accounting", "stats"})
+
+__all__ = [
+    "AccumulationCost",
+    "BudgetContext",
+    "BudgetExhaustedError",
+    "ComponentCost",
+    "FlopscopeArray",
+    "FlopscopeError",
+    "FlopscopeWarning",
+    "NoBudgetContextError",
+    "OpRecord",
+    "PathInfo",
+    "RegimeStep",
+    "StepInfo",
+    "SymmetricTensor",
+    "SymmetryError",
+    "SymmetryGroup",
+    "SymmetryLossWarning",
+    "TimeExhaustedError",
+    "UnsupportedFunctionError",
+    "__numpy_pinned__",
+    "__numpy_supported__",
+    "__numpy_version__",
+    "__version__",
+    "accounting",
+    "as_symmetric",
+    "budget",
+    "budget_live",
+    "budget_reset",
+    "budget_summary",
+    "budget_summary_dict",
+    "clear_cache",
+    "configure",
+    "einsum_accumulation_cost",
+    "einsum_cache_info",
+    "einsum_clear_caches",
+    "is_symmetric",
+    "namespace",
+    "numpy",
+    "reduction_accumulation_cost",
+    "reduction_cache_info",
+    "reduction_clear_cache",
+    "stats",
+    "symmetrize",
+    "tier2_reduction_cost",
+]
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        module = _importlib.import_module(f"flopscope.{name}")
+        globals()[name] = module
+        return module
+    raise AttributeError(
+        f"flopscope does not provide {name!r}. "
+        f"Numpy-shaped operations live under 'flopscope.numpy' "
+        f"(try `import flopscope.numpy as fnp; fnp.{name}`)."
+    )
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | _LAZY_SUBMODULES)

flopscope/_accumulation/__init__.py

@@ -0,0 +1,13 @@
+"""Symmetry-aware einsum accumulation cost — JS-mirrored α/M ladder."""
+
+from ._cost import AccumulationCost, ComponentCost
+from ._ladder import RegimeStep
+from ._public import einsum_accumulation_cost, reduction_accumulation_cost
+
+__all__ = [
+    "AccumulationCost",
+    "ComponentCost",
+    "RegimeStep",
+    "einsum_accumulation_cost",
+    "reduction_accumulation_cost",
+]

flopscope/_accumulation/_bipartite.py

@@ -0,0 +1,121 @@
+"""Bipartite graph + incidence matrix construction for the σ-loop.
+
+(buildBipartite + buildIncidenceMatrix only — runSigmaLoop lives in _detection.py).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class UVertex:
+    op_idx: int
+    class_id: int
+    labels: frozenset[str]
+
+
+@dataclass(frozen=True)
+class BipartiteGraph:
+    u_vertices: tuple[UVertex, ...]
+    incidence: tuple[dict[str, int], ...]  # one per u-vertex
+    u_operand: tuple[int, ...]
+    operand_labels: tuple[frozenset[str], ...]
+    all_labels: tuple[str, ...]  # sorted union
+    free_labels: frozenset[str]
+    summed_labels: frozenset[str]
+    identical_groups: tuple[tuple[int, ...], ...]
+    num_operands: int
+
+
+@dataclass(frozen=True)
+class IncidenceMatrix:
+    matrix: tuple[tuple[int, ...], ...]
+    labels: tuple[str, ...]
+    col_fingerprints: dict[str, tuple[int, ...]]
+    fp_to_labels: dict[tuple[int, ...], frozenset[str]]
+
+
+def build_bipartite(
+    *,
+    subscripts: Sequence[str],
+    output: str,
+    operand_names: Sequence[str],
+) -> BipartiteGraph:
+    """Build the bipartite graph: one U-vertex per axis of each operand.
+
+    Mirrors algorithm.js#buildBipartite. No axis merging — per-operand symmetry
+    is handled later by the σ-loop's wreath enumeration, not by collapsing axes here.
+    """
+    num_ops = len(subscripts)
+    u_vertices: list[UVertex] = []
+    incidence: list[dict[str, int]] = []
+    u_operand: list[int] = []
+    operand_labels: list[frozenset[str]] = []
+
+    for op_idx, sub in enumerate(subscripts):
+        operand_labels.append(frozenset(sub))
+        for axis_idx, ch in enumerate(sub):
+            u_vertices.append(
+                UVertex(op_idx=op_idx, class_id=axis_idx, labels=frozenset({ch}))
+            )
+            incidence.append({ch: 1})
+            u_operand.append(op_idx)
+
+    all_labels_set: set[str] = set()
+    for sub in subscripts:
+        all_labels_set.update(sub)
+    all_labels = tuple(sorted(all_labels_set))
+    output_set = set(output)
+    free_labels = frozenset(lbl for lbl in all_labels if lbl in output_set)
+    summed_labels = frozenset(lbl for lbl in all_labels if lbl not in output_set)
+
+    # Group operand positions by name (Python id() equivalent at the einsum-level
+    # is "same operand_name").
+    name_to_positions: dict[str, list[int]] = {}
+    for i, name in enumerate(operand_names):
+        name_to_positions.setdefault(name, []).append(i)
+    identical_groups = tuple(
+        tuple(positions)
+        for positions in name_to_positions.values()
+        if len(positions) >= 2
+    )
+
+    return BipartiteGraph(
+        u_vertices=tuple(u_vertices),
+        incidence=tuple(incidence),
+        u_operand=tuple(u_operand),
+        operand_labels=tuple(operand_labels),
+        all_labels=all_labels,
+        free_labels=free_labels,
+        summed_labels=summed_labels,
+        identical_groups=identical_groups,
+        num_operands=num_ops,
+    )
+
+
+def build_incidence_matrix(graph: BipartiteGraph) -> IncidenceMatrix:
+    """Build the dense incidence matrix and column fingerprints.
+
+    Mirrors algorithm.js#buildIncidenceMatrix.
+    """
+    labels = graph.all_labels
+    matrix = tuple(
+        tuple(graph.incidence[row_idx].get(lbl, 0) for lbl in labels)
+        for row_idx in range(len(graph.u_vertices))
+    )
+    col_fingerprints: dict[str, tuple[int, ...]] = {}
+    fp_to_labels_mut: dict[tuple[int, ...], set[str]] = {}
+    for c, label in enumerate(labels):
+        fp = tuple(row[c] for row in matrix)
+        col_fingerprints[label] = fp
+        fp_to_labels_mut.setdefault(fp, set()).add(label)
+    fp_to_labels = {fp: frozenset(s) for fp, s in fp_to_labels_mut.items()}
+
+    return IncidenceMatrix(
+        matrix=matrix,
+        labels=labels,
+        col_fingerprints=col_fingerprints,
+        fp_to_labels=fp_to_labels,
+    )

flopscope/_accumulation/_burnside.py

@@ -0,0 +1,51 @@
+"""Size-aware Burnside lemma for orbit counting under heterogeneous label dimensions.
+
+M = (1 / |G|) · Σ_{g ∈ G} ∏_{c ∈ cycles(g)} n_c
+
+where n_c is the common size of the labels in cycle c. Within any cycle of a valid
+symmetry, all labels must share a size (asserted at group-construction time elsewhere;
+re-asserted here for safety).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+
+from flopscope._perm_group import _Permutation as Permutation
+
+
+def _common_size_or_throw(cycle: Sequence[int], sizes: Sequence[int]) -> int:
+    n0 = sizes[cycle[0]]
+    for idx in cycle:
+        if sizes[idx] != n0:
+            cycle_str = ",".join(str(i) for i in cycle)
+            cycle_sizes = ",".join(str(sizes[i]) for i in cycle)
+            raise ValueError(
+                f"cycle size mismatch: labels {cycle_str} have sizes {cycle_sizes} — "
+                f"a permutation can only mix labels of equal size."
+            )
+    return n0
+
+
+def size_aware_burnside(elements: Iterable[Permutation], sizes: Sequence[int]) -> int:
+    """Count orbits of `elements` acting on the assignment grid ∏ [sizes].
+
+    Returns ``M = |X / G|`` where ``X = ∏_ℓ [sizes[ℓ]]``.
+    """
+    elements_tuple = tuple(elements)
+    if not elements_tuple:
+        raise ValueError("size_aware_burnside requires at least one group element")
+
+    total = 0
+    for g in elements_tuple:
+        contribution = 1
+        for cycle in g.full_cyclic_form:
+            contribution *= _common_size_or_throw(cycle, sizes)
+        total += contribution
+
+    if total % len(elements_tuple) != 0:
+        raise ValueError(
+            f"Burnside sum {total} not divisible by |G|={len(elements_tuple)} — "
+            f"group elements probably incomplete or inconsistent."
+        )
+    return total // len(elements_tuple)

flopscope/_accumulation/_cache.py

@@ -0,0 +1,146 @@
+"""LRU cache for compute_accumulation_cost.
+
+Lives in _accumulation rather than _einsum.py so that both the public inspection
+function (_public.einsum_accumulation_cost) and the einsum path (_einsum._get_accumulation_cost)
+can share it without circular imports.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+
+from ._cost import AccumulationCost, compute_accumulation_cost
+
+
+def _make_accumulation_cache(maxsize: int):
+    @functools.lru_cache(maxsize=maxsize)
+    def _compute(
+        canonical_subscripts: str,
+        input_parts: tuple,
+        output_subscript: str,
+        shapes: tuple,
+        sym_fingerprint: tuple,
+        identity_pattern: tuple | None,
+        partition_budget: int | None,
+    ) -> AccumulationCost:
+        # Reconstruct per-op symmetries from the fingerprint.
+        from flopscope._perm_group import SymmetryGroup
+        from flopscope._perm_group import _PermutationCompat as Permutation
+
+        per_op_symmetries: list[Any] = []
+        for fp_entry in sym_fingerprint:
+            if fp_entry is None:
+                per_op_symmetries.append(None)
+                continue
+            axes, gen_arrays = fp_entry
+            gens = [Permutation(list(g)) for g in gen_arrays]
+            group = SymmetryGroup(*gens, axes=axes) if gens else None
+            per_op_symmetries.append(group)
+
+        return compute_accumulation_cost(
+            canonical_subscripts=canonical_subscripts,
+            input_parts=input_parts,
+            output_subscript=output_subscript,
+            shapes=shapes,
+            per_op_symmetries=tuple(per_op_symmetries),
+            identity_pattern=identity_pattern,
+            partition_budget=partition_budget,
+        )
+
+    return _compute
+
+
+_accumulation_cache = _make_accumulation_cache(4096)
+
+
+def get_accumulation_cost_cached(
+    *,
+    canonical_subscripts: str,
+    input_parts: tuple,
+    output_subscript: str,
+    shapes: tuple,
+    sym_fingerprint: tuple,
+    identity_pattern: tuple | None,
+    partition_budget: int | None,
+) -> AccumulationCost:
+    """Cached entry point. Routed through by both public and einsum-internal callers."""
+    return _accumulation_cache(
+        canonical_subscripts,
+        tuple(input_parts),
+        output_subscript,
+        shapes,
+        sym_fingerprint,
+        identity_pattern,
+        partition_budget,
+    )
+
+
+def rebuild_accumulation_cache(maxsize: int) -> None:
+    """Rebuild the cache with a new maxsize."""
+    global _accumulation_cache
+    _accumulation_cache = _make_accumulation_cache(maxsize)
+
+
+def _make_reduction_cache(maxsize: int):
+    @functools.lru_cache(maxsize=maxsize)
+    def _compute(
+        input_shape: tuple,
+        axes_summed: tuple,
+        sym_fingerprint: tuple,
+        op_factor: int,
+        extra_ops: int,
+        partition_budget: int | None,
+    ) -> AccumulationCost:
+        from flopscope._perm_group import SymmetryGroup
+        from flopscope._perm_group import _PermutationCompat as Permutation
+
+        # Reconstruct symmetry from fingerprint.
+        if sym_fingerprint == (None,) or sym_fingerprint == ():
+            symmetry = None
+        else:
+            axes, gen_arrays = sym_fingerprint
+            gens = [Permutation(list(g)) for g in gen_arrays]
+            symmetry = SymmetryGroup(*gens, axes=axes) if gens else None
+
+        from ._reduction import compute_reduction_accumulation_cost
+
+        return compute_reduction_accumulation_cost(
+            input_shape=input_shape,
+            axes_summed=axes_summed,
+            symmetry=symmetry,
+            op_factor=op_factor,
+            extra_ops=extra_ops,
+            partition_budget=partition_budget,
+        )
+
+    return _compute
+
+
+_reduction_cache = _make_reduction_cache(4096)
+
+
+def get_reduction_cost_cached(
+    *,
+    input_shape: tuple,
+    axes_summed: tuple,
+    sym_fingerprint: tuple,
+    op_factor: int,
+    extra_ops: int,
+    partition_budget: int | None,
+) -> AccumulationCost:
+    """Cached entry point for reduction-cost lookups."""
+    return _reduction_cache(
+        input_shape,
+        axes_summed,
+        sym_fingerprint,
+        op_factor,
+        extra_ops,
+        partition_budget,
+    )
+
+
+def rebuild_reduction_cache(maxsize: int) -> None:
+    """Rebuild the cache with a new maxsize."""
+    global _reduction_cache
+    _reduction_cache = _make_reduction_cache(maxsize)