    def _get_max_seq_lens_from_offsets(
        inputs: List[torch.Tensor],
        jagged_offsets_batch_dims: Set[int],
    ) -> Dict[int, int]:
        """
        Get the maximum sequence length encoded in each offsets tensor.

        Offsets tensors encode the length of each sequence in the corresponding
        jagged tensor. Here we extract the maximum sequence length in each offsets
        tensor in the inputs and associate it with the total_length (== offsets[-1])
        of the corresponding jagged tensor in the inputs.
        """
        offsets_inputs = [
            inp
            for inp in inputs
            # offsets tensors are rakn-1 and have a specific first dimension
            if len(inp.shape) == 1 and inp.shape[0] in jagged_offsets_batch_dims
        ]

        max_seq_lens = {}
        for offsets in offsets_inputs:
            offsets = offsets.cpu()
            # the last value in the offsets tensor is the total_length
            # dimension of the corresponding jagged tensor
            total_length = offsets[-1].item()
            # max. sequence length == max. consecutive offset difference
            max_seq_len = torch.max(offsets[1:] - offsets[:-1]).item()
            if total_length in max_seq_lens:
                # if multiple jagged tensors have the same total length,
                # set the max_seq_len to the maximum of all sequences
                # in all corresponding offsets tensors
                max_seq_lens[total_length] = max(
                    max_seq_lens[total_length], max_seq_len
                )
            else:
                max_seq_lens[total_length] = max_seq_len

            logger.info(
                f"Maximum sequence length {max_seq_lens[total_length]} "
                f"for the input jagged tensor with {total_length=} "
                "inferred from the offsets tensor."
            )

        return max_seq_lens
Get the maximum sequence length encoded in each offsets tensor.

Offsets tensors encode the length of each sequence in the corresponding
jagged tensor. Here we extract the maximum sequence length in each offsets
tensor in the inputs and associate it with the total_length (== offsets[-1])
of the corresponding jagged tensor in the inputs.
fx2ait/fx2ait/tensor_spec.py