    def try_getting_jagged_tensor_map(
        cls,
        inputs: List[torch.Tensor],
        jagged_tensor_batch_dims: Set[int],
        fx_inputs: Optional[List[torch.fx.Node]] = None,
    ) -> Optional[Dict[int, int]]:
        """
        Try getting a map associating each jagged tensor input with the offsets.

        In the jagged tensor operators, jagged tensors are normally passed as inputs
        together with the (one or more) offsets tensors. Here we try to infer the
        mapping between each jagged tensor input in the graph and the corresponding
        offsets tensor based on the pre-populated meta-parameters of the fx.Nodes
        corresponding to the jagged tensor inputs.
        """
        if fx_inputs is None or len(fx_inputs) != len(inputs):
            return None

        fx_input_shapes = []
        for fx_inp in fx_inputs:
            tensor_meta = fx_inp.meta.get("tensor_meta", None)
            if tensor_meta is None:
                return None
            shape = getattr(tensor_meta, "shape", None)
            if shape is None:
                return None
            fx_input_shapes.append(shape)

        for i, inp in enumerate(inputs):
            if inp.shape != fx_input_shapes[i]:
                # shape mismatch between the sample input
                # and the corresponding fx.Node in the graph
                return None

        jagged_tensor_map = {}
        seen_offsets_names = {}
        for i, inp in enumerate(inputs):
            if inp.shape[0] in jagged_tensor_batch_dims:
                offsets_name = fx_inputs[i].meta.get("offsets_name", None)
                if offsets_name is None:
                    # not a jagged tensor
                    continue
                if len(offsets_name) > 1:
                    # offsets name attached to the jagged tensor's
                    # fx.Node is either ambiguous: failing here
                    return None
                offsets_name = list(offsets_name)[0]
                if offsets_name not in seen_offsets_names:
                    seen_offsets_names[offsets_name] = len(seen_offsets_names)
                # associate the jagged tensor with the corresponding offsets
                jagged_tensor_map[i] = seen_offsets_names[offsets_name]

        return jagged_tensor_map
Try getting a map associating each jagged tensor input with the offsets.

In the jagged tensor operators, jagged tensors are normally passed as inputs
together with the (one or more) offsets tensors. Here we try to infer the
mapping between each jagged tensor input in the graph and the corresponding
offsets tensor based on the pre-populated meta-parameters of the fx.Nodes
corresponding to the jagged tensor inputs.
fx2ait/fx2ait/tensor_spec.py