@aRI0U Your method is greatly helpful. And I just made a script based on your description, with an additional function to fetch lengths
from typing import Tuple

import torch
from torch import Tensor
from torch import jit
from torch.nn.utils.rnn import pack_sequence, PackedSequence
from torch.nn.utils.rnn import pad_packed_sequence


@jit.script
def sorted_lengths(pack: PackedSequence) -> Tuple[Tensor, Tensor]:
    indices = torch.arange(
        pack.batch_sizes[0],
        dtype=pack.batch_sizes.dtype,
        device=pack.batch_sizes.device,
    )
    lengths = ((indices + 1)[:, None] <= pack.batch_sizes[None, :]).long().sum(dim=1)
    return lengths, indices


@jit.script
def sorted_first_indices(pack: PackedSequence) -> Tensor:
    return torch.arange(
        pack.batch_sizes[0],
        dtype=pack.batch_sizes.dtype,
        device=pack.batch_sizes.device,
    )


@jit.script
def sorted_last_indices(pack: PackedSequence) -> Tensor:
    lengths, indices = sorted_lengths(pack)
    cum_batch_sizes = torch.cat([
        pack.batch_sizes.new_zeros((2,)),
        torch.cumsum(pack.batch_sizes, dim=0),
    ], dim=0)
    return cum_batch_sizes[lengths] + indices


@jit.script
def first_items(pack: PackedSequence, unsort: bool) -> Tensor:
    if unsort and pack.unsorted_indices is not None:
        return pack.data[pack.unsorted_indices]
    else:
        return pack.data[:pack.batch_sizes[0]]


@jit.script
def last_items(pack: PackedSequence, unsort: bool) -> Tensor:
    indices = sorted_last_indices(pack=pack)
    if unsort and pack.unsorted_indices is not None:
        indices = indices[pack.unsorted_indices]
    return pack.data[indices]


if __name__ == '__main__':
    x = pack_sequence([
        torch.randperm(5) + 1,
        torch.randperm(2) + 1,
        torch.randperm(3) + 1,
    ], enforce_sorted=False)
    z, _ = pad_packed_sequence(x, batch_first=True)
    print(z)

    print(first_items(pack=x, unsort=True))
    print(last_items(pack=x, unsort=True))
    print(first_items(pack=x, unsort=False))
    print(last_items(pack=x, unsort=False))

# tensor([[1, 5, 3, 2, 4],
#         [2, 1, 0, 0, 0],
#         [1, 2, 3, 0, 0]])
# tensor([1, 2, 1])
# tensor([4, 1, 3])
# tensor([1, 1, 2])
# tensor([4, 3, 1])