bound_stats

BoundStats

Bases: NumericStats

Base for degree-constraint upper-bound statistics: a guaranteed bound on the number of non-fill values, built from ell_p-norms of degree sequences.

Subclasses differ only in which p values they keep (ps) and how they turn the constraints into an estimate (:meth:estimate_non_fill_values); the scan and all propagation (in :class:BoundStatsFactory) are shared.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class BoundStats(NumericStats):
    """
    Base for degree-constraint upper-bound statistics: a guaranteed bound on the
    number of non-fill values, built from ``ell_p``-norms of degree sequences.

    Subclasses differ only in which ``p`` values they keep (``ps``) and how they
    turn the constraints into an estimate (:meth:`estimate_non_fill_values`); the
    scan and all propagation (in :class:`BoundStatsFactory`) are shared.
    """

    default_ps: ClassVar[tuple[float, ...]] = ()

    def __init__(
        self,
        base: BaseTensorStats,
        dcs: Iterable[DC],
        ps: Iterable[float] | None = None,
    ):
        super().__init__(base)
        self.ps: tuple[float, ...] = tuple(self.default_ps if ps is None else ps)
        self.dcs = set(dcs)

    @abstractmethod
    def estimate_non_fill_values(self) -> float: ...

    def get_embedding(self) -> np.ndarray:
        sizes = [float(self.dim_sizes[field]) for field in self.index_order]
        dc_embedding = [
            dc.value
            for dc in sorted(
                self.dcs,
                key=lambda dc: (
                    tuple(sorted(str(f) for f in dc.from_indices)),
                    tuple(sorted(str(f) for f in dc.to_indices)),
                    dc.p,
                ),
            )
        ]
        return np.log2(np.array([*sizes, *dc_embedding]))

BoundStatsFactory

Bases: BaseTensorStatsFactory[TS], StatsFactory[TS], Generic[TS]

Shared statistics propagation for :class:BoundStats subclasses, combining degree constraints keyed by (from_indices, to_indices, p). The only per-subclass knob is ps.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class BoundStatsFactory(BaseTensorStatsFactory[TS], StatsFactory[TS], Generic[TS]):
    """
    Shared statistics propagation for :class:`BoundStats` subclasses, combining
    degree constraints keyed by ``(from_indices, to_indices, p)``. The only
    per-subclass knob is ``ps``.
    """

    def __init__(self, stats_cls: type[TS], ps: Iterable[float] | None = None):
        super().__init__(stats_cls)
        self.ps: tuple[float, ...] = tuple(stats_cls.default_ps if ps is None else ps)

    def __call__(self, tensor: Any, fields: tuple[Field, ...]) -> TS:
        base = super().__call__(tensor, fields)
        dcs = self.compute_dcs(tensor, fields, self.ps)
        return self.stats_cls(base, dcs, self.ps)

    # Per field: the distinct-value count DC(∅, {i}) and, for each p, the ell_p
    # norm DC({i}, {*fields}, p); plus the full-tensor nnz DC(∅, {*fields}). The
    # ∅-conditioned records are length-1 sequences (p-independent), so kept once.
    @staticmethod
    def compute_dcs(
        tensor: Tensor,
        fields: Iterable[Field],
        ps: Iterable[float],
    ) -> set[DC]:
        fields = list(fields)
        if tensor.ndim == 0:
            return {DC(frozenset(), frozenset(), 1.0)}
        dcs: set[DC] = set()
        norms = get_lp_norms(tensor, fields, ps)
        for field in fields:
            for i, p in enumerate(ps):
                if p == 0:
                    dcs.add(
                        DC(frozenset(), frozenset({field}), norms[field][i], math.inf)
                    )
                elif p == 1:
                    dcs.add(
                        DC(frozenset(), frozenset(fields), norms[field][i], math.inf)
                    )
                else:
                    dcs.add(
                        DC(frozenset({field}), frozenset(fields), norms[field][i], p)
                    )
        return dcs

    def _rebuild(self, base: BaseTensorStats, dcs: Iterable[DC]) -> TS:
        return self.stats_cls(base, dcs, self.ps)

    def _mapjoin_union(self, op: FinchOperator, *union_args: TS) -> TS:
        base = super()._mapjoin_defs(op, *union_args)

        if len(union_args) == 1:
            return self._rebuild(base, union_args[0].dcs)

        dc_keys: Counter[_Key] = Counter()
        stats_dcs: list[dict[_Key, float]] = []
        for stats in union_args:
            dcs: dict[_Key, float] = {}
            extra_fields = tuple(
                x for x in base.index_order if x not in stats.index_order
            )
            extra_size = base.get_dim_space_size(extra_fields)
            for dc in stats.dcs:
                new_key = (dc.from_indices, dc.to_indices, dc.p)
                dcs[new_key] = dc.value
                dc_keys[new_key] += 1

                ext_key = (
                    dc.from_indices,
                    dc.to_indices | frozenset(extra_fields),
                    dc.p,
                )
                if ext_key not in dcs:
                    dc_keys[ext_key] += 1
                prev = dcs.get(ext_key, math.inf)
                dcs[ext_key] = min(prev, dc.value * extra_size)
            stats_dcs.append(dcs)

        new_dcs: dict[_Key, float] = {}
        for key, count in dc_keys.items():
            if count == len(union_args):
                total = sum(d.get(key, 0.0) for d in stats_dcs)
                _, to_indices, _ = key
                if to_indices.issubset(base.index_order):
                    total = min(total, base.get_dim_space_size(to_indices))
                new_dcs[key] = min(float(2**64), total)

        new_stats = {
            DC(from_indices, to_indices, value, p)
            for (from_indices, to_indices, p), value in new_dcs.items()
        }
        return self._rebuild(base, new_stats)

    def _mapjoin_join(self, op: FinchOperator, *join_args: TS) -> TS:
        base = super()._mapjoin_defs(op, *join_args)

        if len(join_args) == 1:
            return self._rebuild(base, join_args[0].dcs)

        new_dc: dict[_Key, float] = {}
        for stats in join_args:
            for dc in stats.dcs:
                key = (dc.from_indices, dc.to_indices, dc.p)
                if dc.value < new_dc.get(key, math.inf):
                    new_dc[key] = dc.value

        new_stats = {
            DC(from_indices, to_indices, value, p)
            for (from_indices, to_indices, p), value in new_dc.items()
        }
        return self._rebuild(base, new_stats)

    def aggregate(
        self,
        op: FinchOperator,
        init: Any | None,
        reduce_indices: tuple[Field, ...],
        stats: TS,
    ) -> TS:
        if len(reduce_indices) == 0:
            base: BaseTensorStats = stats.copy()
        else:
            base = self.aggregate_def(op, init, reduce_indices, stats)
        return self._rebuild(base, stats.dcs)

    def relabel(self, stats: TS, relabel_indices: tuple[Field, ...]) -> TS:
        base = self.relabel_def(stats, relabel_indices)
        mapping = dict(zip(stats.index_order, relabel_indices, strict=True))
        dcs = {
            DC(
                frozenset(mapping.get(f, f) for f in dc.from_indices),
                frozenset(mapping.get(f, f) for f in dc.to_indices),
                dc.value,
                dc.p,
            )
            for dc in stats.dcs
        }
        return self._rebuild(base, dcs)

    def reorder(self, stats: TS, reorder_indices: tuple[Field, ...]) -> TS:
        base = self.reorder_def(stats, reorder_indices)
        return self._rebuild(base, stats.dcs)

DC dataclass

A degree constraint: value bounds ||deg(to_indices | from_indices)||_p.

p = inf (the default) is the classic max-degree constraint; p = 1 is the distinct count; intermediate p capture skew between the two.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

@dataclass(frozen=True)
class DC:
    """
    A degree constraint: ``value`` bounds ``||deg(to_indices | from_indices)||_p``.

    ``p = inf`` (the default) is the classic max-degree constraint; ``p = 1`` is
    the distinct count; intermediate ``p`` capture skew between the two.
    """

    from_indices: frozenset[Field]
    to_indices: frozenset[Field]
    value: float
    p: float = math.inf

DCStats

Bases: BoundStats

Classic degree constraints -- max degrees (p = inf) and distinct counts. Estimates size with a best-first product search over the constraints, clamped by the tensor's dense capacity.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class DCStats(BoundStats):
    """
    Classic degree constraints -- max degrees (``p = inf``) and distinct counts.
    Estimates size with a best-first product search over the constraints,
    clamped by the tensor's dense capacity.
    """

    default_ps: ClassVar[tuple[float, ...]] = (math.inf,)

    def estimate_non_fill_values(self) -> float:
        """Smallest degree-constraint product covering all indices, clamped by
        dense capacity."""
        idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
        if len(idx) == 0:
            return 1.0

        best: dict[frozenset[Field], float] = {frozenset(): 1.0}
        frontier: set[frozenset[Field]] = {frozenset()}

        while True:
            current_bound = best.get(idx, math.inf)
            new_frontier: set[frozenset[Field]] = set()

            for node in frontier:
                for dc in self.dcs:
                    if node.issuperset(dc.from_indices):
                        y = node.union(dc.to_indices)
                        if best[node] > float(2 ** (64 - 2)) or float(dc.value) > float(
                            2 ** (64 - 2)
                        ):
                            y_weight = float(2**64)
                        else:
                            y_weight = best[node] * dc.value
                        if min(current_bound, best.get(y, math.inf)) > y_weight:
                            best[y] = y_weight
                            new_frontier.add(y)
            if len(new_frontier) == 0:
                break
            frontier = new_frontier

        min_weight = float(self.get_dim_space_size(idx))
        for node, weight in best.items():
            if node.issuperset(idx):
                min_weight = min(min_weight, weight)
        return min_weight

estimate_non_fill_values()

Smallest degree-constraint product covering all indices, clamped by dense capacity.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

def estimate_non_fill_values(self) -> float:
    """Smallest degree-constraint product covering all indices, clamped by
    dense capacity."""
    idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
    if len(idx) == 0:
        return 1.0

    best: dict[frozenset[Field], float] = {frozenset(): 1.0}
    frontier: set[frozenset[Field]] = {frozenset()}

    while True:
        current_bound = best.get(idx, math.inf)
        new_frontier: set[frozenset[Field]] = set()

        for node in frontier:
            for dc in self.dcs:
                if node.issuperset(dc.from_indices):
                    y = node.union(dc.to_indices)
                    if best[node] > float(2 ** (64 - 2)) or float(dc.value) > float(
                        2 ** (64 - 2)
                    ):
                        y_weight = float(2**64)
                    else:
                        y_weight = best[node] * dc.value
                    if min(current_bound, best.get(y, math.inf)) > y_weight:
                        best[y] = y_weight
                        new_frontier.add(y)
        if len(new_frontier) == 0:
            break
        frontier = new_frontier

    min_weight = float(self.get_dim_space_size(idx))
    for node, weight in best.items():
        if node.issuperset(idx):
            min_weight = min(min_weight, weight)
    return min_weight

DCStatsFactory

Bases: BoundStatsFactory['DCStats']

Factory for :class:DCStats -- keeps only max-degree (p = inf).

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class DCStatsFactory(BoundStatsFactory["DCStats"]):
    """Factory for :class:`DCStats` -- keeps only max-degree (``p = inf``)."""

    def __init__(self):
        super().__init__(DCStats, ps=(0, 1, math.inf))

LPStats

Bases: BoundStats

LpBound statistics: keeps ell_p-norms of degree sequences for a configurable set of p (not just max degree), and bounds size via the LPflow network-flow LP (:func:_flow_bound).

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class LPStats(BoundStats):
    """
    LpBound statistics: keeps ``ell_p``-norms of degree sequences for a
    configurable set of ``p`` (not just max degree), and bounds size via the
    ``LPflow`` network-flow LP (:func:`_flow_bound`).
    """

    default_ps: ClassVar[tuple[float, ...]] = DEFAULT_PS

    def estimate_non_fill_values(self) -> float:
        idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
        if len(idx) == 0:
            return 1.0

        # An empty (all-fill) input has no non-fill values.
        for dc in self.dcs:
            if len(dc.from_indices) == 0 and dc.value == 0:
                return 0.0

        positive = [dc for dc in self.dcs if dc.value > 0]
        dense_cap = float(self.get_dim_space_size(idx))
        sol = _flow_bound(positive, list(self.dim_sizes.keys()))
        if sol is None:
            return dense_cap
        try:
            bound = math.pow(2.0, sol)
        except OverflowError:
            return dense_cap
        return min(bound, dense_cap)

LPStatsFactory

Bases: BoundStatsFactory['LPStats']

Factory for :class:LPStats with a configurable set of p norms.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class LPStatsFactory(BoundStatsFactory["LPStats"]):
    """Factory for :class:`LPStats` with a configurable set of ``p`` norms."""

    def __init__(self, ps: Iterable[float] | None = None):
        super().__init__(LPStats, ps)