Class TensorRepSetList

backends/vulkan/utils.py:1218–1259

This class is a wrapper around a list of TensorRepSet instances that automatically applies a "broadcasting" mechanism. The broadcasting mechanism allows for a single underlying TensorRepSet to be used for multiple tensors.

class TensorRepSetList:
    """
    This class is a wrapper around a list of TensorRepSet instances that automatically
    applies a "broadcasting" mechanism. The broadcasting mechanism allows for a single
    underlying TensorRepSet to be used for multiple tensors.
    """

    def __init__(
        self,
        tensor_repsets: Union[TensorRepSet, List[TensorRepSet]],
    ):
        self.vals: List[TensorRepSet] = (
            tensor_repsets if isinstance(tensor_repsets, list) else [tensor_repsets]
        )

    def __len__(self):
        return len(self.vals)

    def __getitem__(self, idx: int) -> TensorRepSet:
        if idx > 0 and len(self) == 1:
            return self.vals[0]
        if idx >= len(self.vals):
            return set()
        return self.vals[idx]

    def __setitem__(self, idx: int, val: TensorRepSet) -> None:
        if idx > 0 and len(self.vals) == 1:
            self.vals[0] = val
        else:
            self.vals[idx] = val

    def __str__(self) -> str:
        return f"[{', '.join(str(ts) for ts in self.vals)}]"

    def append(self, val: TensorRepSet) -> None:
        return self.vals.append(val)

    def any_is_empty(self) -> bool:
        if len(self.vals) == 0:
            return True

        return any(tensor_repr.is_empty() for tensor_repr in self.vals)


class OpRepSets: