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 
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 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
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)
    }