    def from_input_list_with_batch_size_jagged_tensor(
        cls,
        inputs: List[torch.Tensor],
        max_batch_size: int,
        max_sequence_length: int,
        jagged_tensor_batch_dims: Set[int],
        jagged_offsets_batch_dims: Set[int],
        additional_inputs: List[torch.Tensor] = None,
        infer_max_seq_lens_from_offsets: bool = False,
        fx_inputs: List[torch.fx.Node] = None,
        jagged_tensor_map: Optional[Dict[int, int]] = None,
    ) -> List["TensorSpec"]:
        """
        Most of the recommendation models will work fine using this function.

        We make an assumption that inferred lowerable subgraph inputs will have
        a single batch dimension with the same max batch size.
        """
        max_seq_lens_from_offsets = {}
        if infer_max_seq_lens_from_offsets:
            max_seq_lens_from_offsets = cls._get_max_seq_lens_from_offsets(
                inputs=inputs,
                jagged_offsets_batch_dims=jagged_offsets_batch_dims,
            )

        result: List = []
        result_unsorted: List = []
        left_inputs: List = []
        left_inputs_ind: List = []
        left_additional_inputs: List = []
        for ind, t in enumerate(inputs):
            batch_dim: int = t.shape[0]
            batch_dim_lower_bound: int = 0
            batch_dim_upper_bound: int = 0
            batch_dim_name: str = ""
            if jagged_tensor_map and ind in jagged_tensor_map:
                batch_dim_lower_bound = 0  # when all sequences are empty
                # if the maximum sequence length for this jagged tensor was not
                # inferred from the offsets, we use the globally configured
                # max_sequence_length (passed as argument to this function)
                max_seq_len = max_seq_lens_from_offsets.get(
                    batch_dim, max_sequence_length
                )
                batch_dim_upper_bound = max_batch_size * max_seq_len
                batch_dim_name = f"batch_size_jagged_tensor_id_{jagged_tensor_map[ind]}"
            elif not jagged_tensor_map and batch_dim in jagged_tensor_batch_dims:
                batch_dim_lower_bound = 0
                max_seq_len = max_seq_lens_from_offsets.get(
                    batch_dim, max_sequence_length
                )
                batch_dim_upper_bound = max_batch_size * max_seq_len
                batch_dim_name = f"batch_size_jagged_tensor_{batch_dim}"
            elif batch_dim in jagged_offsets_batch_dims:
                batch_dim_lower_bound = 2  # prefix 0 + at least one offset
                batch_dim_upper_bound = max_batch_size + 1
                batch_dim_name = f"batch_size_jagged_offsets_{batch_dim}"

            if batch_dim_upper_bound > 0:
                shape: List[IntVar] = []
                for i, d in enumerate(t.shape):
                    if i == 0:
                        shape.append(
                            IntVar(
                                values=[
                                    batch_dim_lower_bound,
                                    batch_dim_upper_bound,
                                ],
                                name=batch_dim_name,
                            )
                        )
                    else:
                        shape.append(IntImm(d))
                result_unsorted.append((ind, TensorSpec(shape, t.dtype)))
            else:
                left_inputs.append(t)
                left_inputs_ind.append(ind)

        if additional_inputs:
            for ind in left_inputs_ind:
                left_additional_inputs.append(additional_inputs[ind])
            input_specs_left = TensorSpec.from_two_input_lists_jagged_tensor(
                left_inputs, left_additional_inputs
            )
            assert len(input_specs_left) == len(
                left_inputs_ind
            ), "Unexpected length for left inputs"

            for index, ind_value in enumerate(left_inputs_ind):
                result_unsorted.append((ind_value, input_specs_left[index]))

        else:
            bs_dim = cls.find_batch_size_dim(left_inputs)
            for index, t in enumerate(left_inputs):
                shape: List[IntVar] = []
                for i, d in enumerate(t.shape):
                    if i == bs_dim[index]:
                        shape.append(IntVar([1, max_batch_size], "batch_size"))
                    else:
                        shape.append(IntImm(d))
                result_unsorted.append(
                    (left_inputs_ind[index], TensorSpec(shape, t.dtype))
                )
        result = sorted(result_unsorted, key=lambda num: num[0])
        result = [r[1] for r in result]
        return result
Most of the recommendation models will work fine using this function.

We make an assumption that inferred lowerable subgraph inputs will have
a single batch dimension with the same max batch size.
fx2ait/fx2ait/tensor_spec.py