Class TensorRepSet

backends/vulkan/utils.py:865–1091

This class describes the possible set of representations (i.e. TensorRepr) that may be used to represent a tensor. This set is determined by the implementation of the operator that the tensor participates in as well as the texture extents of the GPU.

class TensorRepSet:
    """
    This class describes the possible set of representations (i.e. TensorRepr) that may
    be used to represent a tensor. This set is determined by the implementation of the
    operator that the tensor participates in as well as the texture extents of the GPU.
    """

    def __init__(
        self,
        buffer_memory_layouts: Set[VkMemoryLayout],
        texture_memory_layouts: Set[VkMemoryLayout],
    ):
        self.valid_buffer_layouts = buffer_memory_layouts
        self.valid_texture_layouts = texture_memory_layouts

    def __str__(self) -> str:
        buffer_layouts = ", ".join(layout.name for layout in self.valid_buffer_layouts)
        texture_layouts = ", ".join(
            layout.name for layout in self.valid_texture_layouts
        )
        return f"TensorRepSet(Buffer Layouts: [{buffer_layouts}], Texture Layouts: [{texture_layouts}])"

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, TensorRepSet):
            return NotImplemented
        return (
            self.valid_buffer_layouts == other.valid_buffer_layouts
            and self.valid_texture_layouts == other.valid_texture_layouts
        )

    def __ne__(self, other: object) -> bool:
        return not self.__eq__(other)

    def copy(self) -> "TensorRepSet":
        return TensorRepSet(
            set(self.valid_buffer_layouts), set(self.valid_texture_layouts)
        )

    def is_empty(self) -> bool:
        """
        A TensorRepSet is "empty" if there are no valid representations of the tensor.
        """
        return (
            len(self.valid_buffer_layouts) == 0 and len(self.valid_texture_layouts) == 0
        )

    def make_intersect(self, other: "TensorRepSet") -> "TensorRepSet":
        """
        Merge this TensorRepr with another TensorRepr, returning a new TensorRepr
        with the intersection of the two.
        """
        return TensorRepSet(
            self.valid_buffer_layouts & other.valid_buffer_layouts,
            self.valid_texture_layouts & other.valid_texture_layouts,
        )

    def make_union(self, other: "TensorRepSet") -> "TensorRepSet":
        """