Skip to content

util

Shared functionality across TensorStats implementations.

get_lp_norms(arr, fields, norms)

Scan arr and return the Lp norms of each axis' degree sequence.

The degree sequence of axis i counts, for each index of that axis, how many non-fill entries carry it. Its norms are what :class:DC degree constraints are built from:

  • p == 0 -- the number of indices with at least one non-fill entry, i.e. the size of the projection onto that axis.
  • p == inf -- the largest degree.
  • otherwise -- (sum(degree ** p)) ** (1 / p).

Parameters:

Name Type Description Default
arr Tensor

The tensor to scan.

 required
fields Iterable[Field]

The axis names of arr (one per dimension), in order.

 required
norms Iterable[float]

The Lp norms to compute.

 required

Returns:

Type Description
dict[Field, list[float]]

A dictionary mapping each field to its norms, ordered as norms.
Source code in src/finch/autoschedule/tensor_stats/util.py 
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def get_lp_norms(
    arr: Tensor, fields: Iterable[Field], norms: Iterable[float]
) -> dict[Field, list[float]]:
    """Scan ``arr`` and return the Lp norms of each axis' degree sequence.

    The degree sequence of axis ``i`` counts, for each index of that axis, how
    many non-fill entries carry it. Its norms are what :class:`DC` degree
    constraints are built from:

    * ``p == 0`` -- the number of indices with at least one non-fill entry,
      i.e. the size of the projection onto that axis.
    * ``p == inf`` -- the largest degree.
    * otherwise -- ``(sum(degree ** p)) ** (1 / p)``.

    Args:
        arr: The tensor to scan.
        fields: The axis names of ``arr`` (one per dimension), in order.
        norms: The Lp norms to compute.

    Returns:
        A dictionary mapping each field to its norms, ordered as ``norms``.
    """
    from finch.autoschedule.default_schedulers import NON_RECURSIVE_SCHEDULER

    fields = tuple(fields)
    norms = tuple(norms)
    if not fields or not norms:
        return {field: [] for field in fields}

    int_zero = Literal(np.int64(0))
    float_zero = Literal(np.float64(0.0))
    # Indicator of the non-fill structure. Left inlined in each reduction below
    # rather than bound to its own alias, so no full-size temporary is built.
    non_fill = MapJoin(
        Literal(ffuncs.ne),
        (Table(Literal(arr), fields), Literal(arr.fill_value)),
    )

    bodies: list[Query] = []
    outputs: list[Alias] = []
    for dim, field in enumerate(fields):
        degrees = Alias(f"degrees_{dim}")
        bodies.append(
            Query(
                Table(degrees, (field,)),
                Aggregate(
                    Literal(ffuncs.add),
                    int_zero,
                    non_fill,
                    tuple(f for f in fields if f != field),
                ),
            )
        )
        degree_table = Table(degrees, (field,))

        for k, norm in enumerate(norms):
            out = Alias(f"norm_{dim}_{k}")
            rhs: Aggregate
            if norm == 0:
                rhs = Aggregate(
                    Literal(ffuncs.add),
                    int_zero,
                    MapJoin(Literal(ffuncs.ne), (degree_table, Literal(np.int64(0)))),
                    (field,),
                )
            elif np.isinf(norm):
                # Degrees are non-negative, so 0 is a safe identity for max.
                rhs = Aggregate(Literal(ffuncs.max), int_zero, degree_table, (field,))
            else:
                # The outer ** (1 / norm) is applied to the scalar result below.
                rhs = Aggregate(
                    Literal(ffuncs.add),
                    float_zero,
                    MapJoin(
                        Literal(ffuncs.pow), (degree_table, Literal(np.float64(norm)))
                    ),
                    (field,),
                )
            bodies.append(Query(Table(out, ()), rhs))
            outputs.append(out)

    prgm = Plan((*bodies, Produces(tuple(outputs))))
    results = NON_RECURSIVE_SCHEDULER(prgm)

    flat: list[float] = []
    all_norms = [norm for _ in fields for norm in norms]
    for result, norm in zip(results, all_norms, strict=True):
        value = float(np.asarray(result)[()])
        if norm != 0 and not np.isinf(norm):
            value = value ** (1.0 / norm)
        flat.append(value)

    return {
        field: flat[i * len(norms) : (i + 1) * len(norms)]
        for i, field in enumerate(fields)
    }