    def _size(self, dim) -> IntVar:
        """
        Gets the size of tensor at dim=dim.
        dim must be between [-rank, rank - 1].
        It should not be used directly in IR, use ops.size(dim) instead.
        """
        return self._attrs["shape"][wrap_dim(dim, self._rank())]
Gets the size of tensor at dim=dim.
dim must be between [-rank, rank - 1].
It should not be used directly in IR, use ops.size(dim) instead.
python/aitemplate/compiler/base.py