Skip to content

tensor_stats

BaseTensorStats

Bases: TensorStats

Base class for tensor statistics.

Source code in src/finch/autoschedule/tensor_stats/tensor_stats.py 
 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
121
122
123
124
class BaseTensorStats(TensorStats):
    """
    Base class for tensor statistics.
    """

    def __init__(
        self,
        index_order: Iterable[Field] | BaseTensorStats,
        dim_sizes: MutableMapping[Field, float] | None = None,
        fill_value: Any = None,
    ):
        if isinstance(index_order, BaseTensorStats):
            if dim_sizes is not None:
                raise TypeError(
                    "dim_sizes must be omitted when constructing from BaseTensorStats"
                )
            self._init_from_fields(
                index_order.index_order, index_order.dim_sizes, index_order.fill_value
            )
        else:
            if dim_sizes is None:
                raise TypeError("dim_sizes are required when constructing from fields")
            self._init_from_fields(index_order, dim_sizes, fill_value)

    def _init_from_fields(
        self,
        index_order: Iterable[Field],
        dim_sizes: MutableMapping[Field, float],
        fill_value: Any,
    ):
        self._index_order = tuple(index_order)
        self._dim_sizes: OrderedDict[Field, float] = OrderedDict(dim_sizes)
        # Stats always carry an AbstractFill: propagation must be able to say
        # "this fill exists but must not be specialized on".
        self._fill_value: AbstractFill = as_fill(fill_value)

    @classmethod
    def from_fields(
        cls,
        index_order: Iterable[Field],
        dim_sizes: MutableMapping[Field, float],
        fill_value: Any,
    ) -> Self:
        return cls(index_order, dim_sizes, fill_value)

    def copy(self) -> Self:
        new = object.__new__(type(self))
        new.__dict__ = {name: copy.copy(value) for name, value in self.__dict__.items()}
        return new

    @property
    def dim_sizes(self) -> MutableMapping[Field, float]:
        return self._dim_sizes

    @dim_sizes.setter
    def dim_sizes(self, value: MutableMapping[Field, float]):
        self._dim_sizes = OrderedDict(value)

    def get_dim_size(self, idx: Field) -> float:
        return self.dim_sizes[idx]

    @property
    def index_order(self) -> tuple[Field, ...]:
        return self._index_order

    @index_order.setter
    def index_order(self, value: Iterable[Field]):
        self._index_order = tuple(value)

    @property
    def idxs(self) -> tuple[Field, ...]:
        return self.index_order

    @property
    def fill_value(self) -> AbstractFill:
        """The fill, and whether kernels may specialize on it."""
        return self._fill_value

    @fill_value.setter
    def fill_value(self, value: Any):
        self._fill_value = as_fill(value)

    def get_dim_space_size(self, idx: Iterable[Field]) -> float:
        prod = 1
        for i in idx:
            prod *= int(self.dim_sizes[i])
            if prod == 0 or prod > np.iinfo(np.int64).max:
                return float("inf")
        return float(prod)

fill_value property writable

The fill, and whether kernels may specialize on it.