flopscope 0.2.0__py3-none-any.whl

flopscope/_accumulation/_reduction.py

@@ -0,0 +1,310 @@
+"""Reduction-cost orchestrator + helpers.
+
+Implements the two-tier reduction cost model:
+- Tier 1 (ufunc.reduce): orbit-mapping α via compute_accumulation_cost
+- Tier 2 (non-ufunc):   dense × (num_output_orbits / num_output_elems)
+
+See .aicrowd/superpowers/specs/2026-05-13-symmetry-aware-reduction-cost-design.md.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from flopscope._perm_group import SymmetryGroup
+
+
+def _normalize_axis(
+    axis: int | tuple[int, ...] | list[int] | None,
+    ndim: int,
+) -> tuple[int, ...]:
+    """Normalize an axis specifier to a sorted tuple of non-negative ints.
+
+    - None → all axes (full reduction)
+    - int → singleton, with negative-index wrapping
+    - tuple/list → sorted unique, with negative-index wrapping
+    """
+    if axis is None or ndim == 0:
+        # 0-d input has no reducible axes. Numpy accepts ``axis=0`` on 0-d
+        # arrays as a no-op (matching ``ufunc.reduce``'s default) and raises
+        # for any other concrete axis before this cost path is reached, so
+        # we short-circuit instead of evaluating ``axis % 0``.
+        return tuple(range(ndim))
+    if isinstance(axis, int):
+        return (axis % ndim,)
+    return tuple(sorted({a % ndim for a in axis}))
+
+
+def _num_output_orbits(
+    input_shape: tuple[int, ...],
+    axes_summed: tuple[int, ...],
+    symmetry: SymmetryGroup | None,
+) -> int:
+    """Count orbits of the output multi-index under the output stabilizer.
+
+    Output shape = input_shape with axes_summed removed.
+    Output stabilizer = reduce_group(symmetry, ndim=ndim, axis=axes_summed).
+    Returns the size-aware Burnside orbit count on the output shape.
+    """
+    ndim = len(input_shape)
+    axes_set = frozenset(axes_summed)
+    output_shape = tuple(input_shape[i] for i in range(ndim) if i not in axes_set)
+
+    if not output_shape:
+        return 1
+
+    if symmetry is None:
+        return math.prod(output_shape)
+
+    from flopscope._symmetry_utils import reduce_group
+
+    from ._burnside import size_aware_burnside
+
+    output_group = reduce_group(
+        symmetry,
+        ndim=ndim,
+        axis=axes_summed,
+        keepdims=False,
+    )
+
+    if output_group is None:
+        return math.prod(output_shape)
+
+    # The group's permutation indices are LOCAL (0..degree-1).
+    # Map them to the axes the group actually operates on in output_shape,
+    # then extract only those sizes so Burnside sees the right array shape.
+    group_output_axes = (
+        output_group.axes
+        if output_group.axes is not None
+        else tuple(range(output_group.degree))
+    )
+    group_sizes = tuple(output_shape[ax] for ax in group_output_axes)
+
+    elements = output_group.elements()
+    orbits_on_group_dims = size_aware_burnside(elements, group_sizes)
+
+    # Remaining output dims (not in the group) contribute independently.
+    group_axes_set = frozenset(group_output_axes)
+    remaining = math.prod(
+        output_shape[i] for i in range(len(output_shape)) if i not in group_axes_set
+    )
+    return orbits_on_group_dims * max(remaining, 1)
+
+
+def output_discounted_reduction_cost(
+    input_shape: tuple[int, ...],
+    axes_summed: tuple[int, ...],
+    symmetry: SymmetryGroup | None,
+    dense_per_output_cost: int,
+) -> int:
+    """Tier-2 reduction cost: num_output_orbits × dense_per_output_cost.
+
+    For non-ufunc reductions (median, percentile, quantile) where we can't
+    decompose the operation internally. The discount uses only the output
+    stabilizer; input symmetry doesn't help unless the operation is
+    decomposable.
+
+    Equivalent to: dense_total × (num_output_orbits / num_output_elems).
+    """
+    num_orbits = _num_output_orbits(input_shape, axes_summed, symmetry)
+    return num_orbits * dense_per_output_cost
+
+
+def compute_reduction_accumulation_cost(
+    input_shape: tuple[int, ...],
+    axes_summed: tuple[int, ...],
+    symmetry: SymmetryGroup | None = None,
+    *,
+    op_factor: int = 1,
+    extra_ops: int = 0,
+    partition_budget: int | None = None,
+):
+    """Tier-1 reduction cost: orbit-mapping α via the existing ladder.
+
+    Constructs a single-operand 'fake einsum' (input subscripts = full label
+    set, output subscripts = labels not in axes_summed), runs the existing
+    component decomposition + ladder, then aggregates via aggregate_reduction.
+
+    Off-by-one fix (#56): cost = op_factor × (α - num_output_orbits) + extra_ops.
+
+    Parameters
+    ----------
+    input_shape : tuple of int
+        Shape of the input array.
+    axes_summed : tuple of int
+        Axes to reduce. Must be non-negative, sorted, unique.
+        Use _normalize_axis to prepare.
+    symmetry : SymmetryGroup, optional
+        Pointwise symmetry of the input. None for dense inputs.
+    op_factor : int, default 1
+        FLOPs per binary accumulation event.
+    extra_ops : int, default 0
+        Output-side dense additions (e.g. num_output_orbits for mean's divide).
+    partition_budget : int, optional
+        Override the partition counter budget. Defaults to global setting.
+
+    Returns
+    -------
+    AccumulationCost
+        Same shape as einsum_accumulation_cost — fields total/mu/alpha/m_total/...
+    """
+    from ._cost import aggregate_reduction, compute_accumulation_cost
+
+    ndim = len(input_shape)
+    if not axes_summed:
+        # No reduction at all — cost is 0 (just identity).
+        return _trivial_zero_cost(input_shape, op_factor, extra_ops)
+
+    # Build a fake einsum: 'abc...' input, output is labels of kept axes.
+    if ndim > 26:
+        # Fall back to dense — labelers can't go beyond a-z cleanly. Unusual.
+        return _dense_fallback_cost(
+            input_shape,
+            axes_summed,
+            symmetry,
+            op_factor,
+            extra_ops,
+        )
+
+    axes_set = frozenset(axes_summed)
+    labels = [chr(ord("a") + i) for i in range(ndim)]
+    input_subs = "".join(labels)
+    output_subs = "".join(labels[i] for i in range(ndim) if i not in axes_set)
+    canonical_subscripts = f"{input_subs}->{output_subs}"
+
+    # Single operand, per-op symmetry = the input symmetry.
+    per_op_symmetries = (symmetry,)
+
+    from flopscope._perm_group import _DiminoBudgetExceeded
+
+    try:
+        einsum_cost = compute_accumulation_cost(
+            canonical_subscripts=canonical_subscripts,
+            input_parts=(input_subs,),
+            output_subscript=output_subs,
+            shapes=(input_shape,),
+            per_op_symmetries=per_op_symmetries,
+            identity_pattern=None,  # single operand, no identity grouping
+            partition_budget=partition_budget,
+        )
+
+        # `output_dense` carries num_output_orbits — orbit-aware, computed
+        # directly via _num_output_orbits. For dense inputs this equals
+        # prod(output_shape) so the locked-stub docstring's intent is preserved.
+        num_output_orbits = _num_output_orbits(
+            input_shape=input_shape,
+            axes_summed=axes_summed,
+            symmetry=symmetry,
+        )
+    except _DiminoBudgetExceeded as exc:
+        # _num_output_orbits' element enumeration can also exceed the budget
+        # on an oversized output stabilizer. Charge dense and warn.
+        # (compute_accumulation_cost emits its own CostFallbackWarning when it
+        # bails before this; if it succeeded but _num_output_orbits bailed,
+        # this is the only warning the caller sees.)
+        import warnings as _warnings
+
+        from flopscope.errors import CostFallbackWarning
+
+        _warnings.warn(
+            CostFallbackWarning(
+                f"reduction: dimino_budget exceeded ({exc.seen_count} > {exc.budget}) "
+                f"while computing output orbit count — charging dense cost. "
+                f"Raise via flopscope.configure(dimino_budget=...) to attempt "
+                f"exact counting."
+            ),
+            stacklevel=3,
+        )
+        return _dense_fallback_cost(
+            input_shape, axes_summed, symmetry, op_factor, extra_ops
+        )
+    dense_baseline = math.prod(input_shape) if input_shape else 1
+
+    # If compute_accumulation_cost already bailed (e.g. dimino_budget exceeded
+    # on an auto-inferred S_n), it returns an AccumulationCost with
+    # fallback_used=True and an empty per_component tuple. aggregate_reduction
+    # would then mis-interpret the empty tuple as "no failures"; re-shape the
+    # cost under the reduction-fallback formula instead.
+    if einsum_cost.fallback_used and not einsum_cost.per_component:
+        from ._cost import AccumulationCost
+
+        input_axis_size = (
+            dense_baseline // num_output_orbits if num_output_orbits else 0
+        )
+        fallback_total = (
+            num_output_orbits * max(0, input_axis_size - 1) * op_factor + extra_ops
+        )
+        return AccumulationCost(
+            total=fallback_total,
+            mu=None,
+            alpha=None,
+            m_total=1,
+            dense_baseline=dense_baseline,
+            num_terms=1,
+            per_component=(),
+            fallback_used=True,
+            unavailable_components=(),
+            unavailable_reason=einsum_cost.unavailable_reason,
+        )
+
+    # Re-aggregate the per-component costs under the reduction formula.
+    return aggregate_reduction(
+        einsum_cost.per_component,
+        op_factor=op_factor,
+        dense_baseline=dense_baseline,
+        output_dense=num_output_orbits,
+        extra_ops=extra_ops,
+    )
+
+
+def _trivial_zero_cost(
+    input_shape: tuple[int, ...],
+    op_factor: int,
+    extra_ops: int,
+):
+    """Cost of a 'no axes reduced' degenerate case."""
+    from ._cost import AccumulationCost
+
+    return AccumulationCost(
+        total=extra_ops,
+        mu=0,
+        alpha=0,
+        m_total=1,
+        dense_baseline=math.prod(input_shape) if input_shape else 1,
+        num_terms=1,
+        per_component=(),
+        fallback_used=False,
+    )
+
+
+def _dense_fallback_cost(
+    input_shape: tuple[int, ...],
+    axes_summed: tuple[int, ...],
+    symmetry: SymmetryGroup | None,
+    op_factor: int,
+    extra_ops: int,
+):
+    """ndim > 26 fallback: charge dense without symmetry."""
+    from ._cost import AccumulationCost
+
+    axes_set = frozenset(axes_summed)
+    dense = math.prod(input_shape) if input_shape else 1
+    output_dense = (
+        math.prod(input_shape[i] for i in range(len(input_shape)) if i not in axes_set)
+        if any(i not in axes_set for i in range(len(input_shape)))
+        else 1
+    )
+    input_axis_size = dense // output_dense if output_dense else 0
+    total = output_dense * max(0, input_axis_size - 1) * op_factor + extra_ops
+    return AccumulationCost(
+        total=total,
+        mu=None,
+        alpha=None,
+        m_total=1,
+        dense_baseline=dense,
+        num_terms=1,
+        per_component=(),
+        fallback_used=True,
+    )

flopscope/_accumulation/_regimes.py

@@ -0,0 +1,303 @@
+"""The four mixed-shape accumulation regimes plus functionalProjection.
+
+Each regime is a Regime instance with `recognize(ctx) -> Verdict` and
+`compute(ctx) -> RegimeOutput`. The dispatcher in _ladder.py iterates them
+in priority order; functionalProjection is checked first (Stage 2a) and the
+remaining three form the mixed-shape ladder (Stage 2b).
+"""
+
+from __future__ import annotations
+
+from ._burnside import size_aware_burnside
+from ._ladder import Regime, RegimeContext, RegimeOutput, Verdict
+from ._output_orbit import projection_is_functional
+
+# ── functionalProjection ─────────────────────────────────────────────
+
+
+def _functional_projection_recognize(ctx: RegimeContext) -> Verdict:
+    if projection_is_functional(ctx.elements, ctx.visible_positions):
+        return Verdict(
+            fired=True,
+            reason="each product orbit reaches exactly one stored output representative",
+        )
+    return Verdict(
+        fired=False,
+        reason="some pointwise symmetry moves an output label into a summed label",
+    )
+
+
+def _functional_projection_compute(ctx: RegimeContext) -> RegimeOutput:
+    count = size_aware_burnside(ctx.elements, ctx.sizes)
+    return RegimeOutput(
+        count=count,
+        sub_steps=(
+            {
+                "step": "projection-functional",
+                "reason": "G preserves V setwise; projection descends to output reps",
+                "count": count,
+            },
+        ),
+    )
+
+
+FUNCTIONAL_PROJECTION_REGIME: Regime = Regime(
+    id="functionalProjection",
+    recognize=_functional_projection_recognize,
+    compute=_functional_projection_compute,
+)
+
+
+# ── singleton (B.5: |V| = 1) ─────────────────────────────────────────
+
+
+def _label_orbit(elements, label_idx: int) -> list[int]:
+    """G-orbit of a single label position."""
+    seen = {label_idx}
+    changed = True
+    while changed:
+        changed = False
+        for g in elements:
+            for p in list(seen):
+                q = g.array_form[p]
+                if q not in seen:
+                    seen.add(q)
+                    changed = True
+    return sorted(seen)
+
+
+def _cycles_on_subset(perm, subset: list[int]) -> int:
+    """Count cycles of perm restricted to a perm-invariant subset."""
+    subset_set = set(subset)
+    seen: set[int] = set()
+    cycles = 0
+    for start in subset:
+        if start in seen:
+            continue
+        cycles += 1
+        cur = start
+        while cur not in seen:
+            if cur not in subset_set:
+                raise ValueError("subset not invariant under perm")
+            seen.add(cur)
+            cur = perm.array_form[cur]
+    return cycles
+
+
+def _subset_cycle_product(perm, subset: list[int], sizes) -> int:
+    """∏ n_c over cycles of perm in `subset`. Each cycle's labels must share a size."""
+    subset_set = set(subset)
+    seen: set[int] = set()
+    product = 1
+    for start in subset:
+        if start in seen:
+            continue
+        cycle: list[int] = []
+        cur = start
+        while cur not in seen:
+            if cur not in subset_set:
+                raise ValueError("subset not invariant under perm")
+            seen.add(cur)
+            cycle.append(cur)
+            cur = perm.array_form[cur]
+        n0 = sizes[cycle[0]]
+        for i in cycle:
+            if sizes[i] != n0:
+                raise ValueError("singleton: cycle in R has mixed sizes")
+        product *= n0
+    return product
+
+
+def _singleton_recognize(ctx: RegimeContext) -> Verdict:
+    if len(ctx.va) == 1:
+        return Verdict(fired=True, reason="|V| = 1")
+    return Verdict(fired=False, reason=f"|V| = {len(ctx.va)}, not 1")
+
+
+def _singleton_compute(ctx: RegimeContext) -> RegimeOutput:
+    v_pos = ctx.visible_positions[0]
+    omega = _label_orbit(ctx.elements, v_pos)
+    n_omega = ctx.sizes[v_pos]
+    for idx in omega:
+        if ctx.sizes[idx] != n_omega:
+            raise ValueError(
+                f"singleton: orbit of label has mixed sizes at {ctx.labels[idx]}"
+            )
+    omega_set = set(omega)
+    rest = [i for i in range(len(ctx.labels)) if i not in omega_set]
+
+    total = 0
+    for g in ctx.elements:
+        rest_factor = _subset_cycle_product(g, rest, ctx.sizes)
+        c_omega = _cycles_on_subset(g, omega)
+        total += rest_factor * (n_omega**c_omega - (n_omega - 1) ** c_omega)
+    count = (n_omega * total) // len(ctx.elements)
+    return RegimeOutput(count=count, sub_steps=())
+
+
+SINGLETON_REGIME: Regime = Regime(
+    id="singleton",
+    recognize=_singleton_recognize,
+    compute=_singleton_compute,
+)
+
+
+# ── young (B.6: G = Sym(L), uniform sizes, |V| ≥ 2) ──────────────────
+
+
+import math
+
+
+def _multiset_count(n: int, k: int) -> int:
+    """Number of size-k multisets from [n]. C(n + k - 1, k)."""
+    if k == 0:
+        return 1
+    num = 1
+    den = 1
+    for i in range(k):
+        num *= n + k - 1 - i
+        den *= i + 1
+    return num // den
+
+
+def _young_recognize(ctx: RegimeContext) -> Verdict:
+    if not ctx.elements or len(ctx.elements) <= 1:
+        return Verdict(fired=False, reason="|G| <= 1")
+    if len(ctx.va) < 2:
+        return Verdict(fired=False, reason="|V| < 2; singleton handles this")
+
+    expected_full_sym = math.factorial(len(ctx.labels))
+    if len(ctx.elements) != expected_full_sym:
+        return Verdict(
+            fired=False,
+            reason=f"|G|={len(ctx.elements)} != |L|!={expected_full_sym}",
+        )
+
+    label_to_idx = {lbl: i for i, lbl in enumerate(ctx.labels)}
+    v_idx_set = {label_to_idx[lbl] for lbl in ctx.va}
+    has_cross = any(
+        any(g.array_form[label_to_idx[lbl]] not in v_idx_set for lbl in ctx.va)
+        for g in ctx.elements
+    )
+    if not has_cross:
+        return Verdict(fired=False, reason="no cross-V/W element")
+
+    if not ctx.sizes:
+        return Verdict(fired=False, reason="no sizes provided")
+
+    n_l = ctx.sizes[0]
+    if any(s != n_l for s in ctx.sizes):
+        return Verdict(fired=False, reason="mixed label sizes")
+
+    return Verdict(fired=True, reason="G = Sym(L); Young equation applies")
+
+
+def _young_compute(ctx: RegimeContext) -> RegimeOutput:
+    n_l = ctx.sizes[0]
+    visible_multisets = _multiset_count(n_l, len(ctx.va))
+    summed_multisets = _multiset_count(n_l, len(ctx.wa))
+    count = visible_multisets * summed_multisets
+    return RegimeOutput(
+        count=count,
+        sub_steps=(
+            {
+                "step": "full-symmetric-output-orbit-formula",
+                "n": n_l,
+                "v_count": len(ctx.va),
+                "w_count": len(ctx.wa),
+                "visible_multisets": visible_multisets,
+                "summed_multisets": summed_multisets,
+                "count": count,
+            },
+        ),
+    )
+
+
+YOUNG_REGIME: Regime = Regime(
+    id="young",
+    recognize=_young_recognize,
+    compute=_young_compute,
+)
+
+
+# ── partitionCount (B.7: typed equality-pattern enumeration) ─────────
+
+
+from ._output_orbit import restrict_stabilizer_to_positions
+from ._partition import (
+    count_map_orbits_under_h,
+    generate_typed_set_partitions,
+    induced_block_action_size,
+    induced_prefix_maps,
+    num_blocks,
+    partition_key,
+    partition_orbit_reps,
+    typed_labeling_count,
+)
+
+
+def _partition_count_recognize(ctx: RegimeContext) -> Verdict:
+    partitions = generate_typed_set_partitions(ctx.sizes)
+    if len(partitions) > ctx.partition_budget:
+        return Verdict(
+            fired=False,
+            reason=(
+                f"typed partition count {len(partitions)} exceeds "
+                f"budget {ctx.partition_budget}"
+            ),
+        )
+    return Verdict(
+        fired=True,
+        reason=f"typed partition count over {len(partitions)} equality patterns",
+    )
+
+
+def _partition_count_compute(ctx: RegimeContext) -> RegimeOutput:
+    h_elements = restrict_stabilizer_to_positions(ctx.elements, ctx.visible_positions)
+    partitions = generate_typed_set_partitions(ctx.sizes)
+    reps = partition_orbit_reps(partitions, ctx.elements)
+
+    total = 0
+    sub_steps: list[dict] = []
+    for partition in reps:
+        labelings = typed_labeling_count(partition, ctx.sizes)
+        block_action = induced_block_action_size(partition, ctx.elements)
+        if labelings % block_action != 0:
+            raise ValueError(
+                f"partition {partition_key(partition)} has labelings={labelings} "
+                f"not divisible by block action size={block_action} — "
+                f"invariant violation in G"
+            )
+        input_orbits = labelings // block_action
+        maps = induced_prefix_maps(partition, ctx.elements, ctx.visible_positions)
+        output_orbits = count_map_orbits_under_h(maps, h_elements)
+        term = input_orbits * output_orbits
+        total += term
+        sub_steps.append(
+            {
+                "partition_key": partition_key(partition),
+                "blocks": num_blocks(partition),
+                "typed_labelings": labelings,
+                "block_action_size": block_action,
+                "input_orbit_count": input_orbits,
+                "induced_map_count": len(maps),
+                "output_orbit_count": output_orbits,
+                "contribution": term,
+            }
+        )
+
+    return RegimeOutput(count=total, sub_steps=tuple(sub_steps))
+
+
+PARTITION_COUNT_REGIME: Regime = Regime(
+    id="partitionCount",
+    recognize=_partition_count_recognize,
+    compute=_partition_count_compute,
+)
+
+
+MIXED_REGIMES: tuple[Regime, ...] = (
+    SINGLETON_REGIME,
+    YOUNG_REGIME,
+    PARTITION_COUNT_REGIME,
+)

flopscope/_accumulation/_shape.py

@@ -0,0 +1,33 @@
+"""Shape classifier — Stage 1 of the cost ladder.
+
+The shape classifier is structural and runs before the regime ladder. It produces
+a label that goes onto the per-component output for diagnostic display alongside
+the regime ID. (regime_id is computational; shape is structural.)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Literal
+
+from flopscope._perm_group import _Permutation as Permutation
+
+Shape = Literal["trivial", "allVisible", "allSummed", "mixed"]
+
+SHAPES: tuple[Shape, ...] = ("trivial", "allVisible", "allSummed", "mixed")
+
+
+def detect_shape(
+    *,
+    va: Sequence[str],
+    wa: Sequence[str],
+    elements: Sequence[Permutation],
+) -> Shape:
+    """Classify a component's structural shape from its V/W partition and group size."""
+    if not elements or len(elements) <= 1:
+        return "trivial"
+    if len(wa) == 0:
+        return "allVisible"
+    if len(va) == 0:
+        return "allSummed"
+    return "mixed"