Class TensorReprList

backends/vulkan/utils.py:809–862

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

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

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

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

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

    def __setitem__(self, idx: int, val: TensorRepr) -> None:
        if idx > 0 and len(self) == 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 __eq__(self, other: object) -> bool:
        if not isinstance(other, TensorReprList):
            return NotImplemented

        if len(self) == len(other):
            for self_val, other_val in zip(self.vals, other.vals):
                if self_val != other_val:
                    return False

            return True

        return False

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

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

    def storage_type(self, idx: int = 0) -> VkStorageType:
        return self.vals[idx].storage_type

    def memory_layout(self, idx: int = 0) -> VkMemoryLayout:
        return self.vals[idx].memory_layout


class TensorRepSet: