onnx-diagnostic 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@ Investigates onnx models.
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.4.0"
+__version__ = "0.4.2"
 __author__ = "Xavier Dupré"

onnx_diagnostic/_command_lines_parser.py

@@ -214,6 +214,22 @@ def get_parser_config() -> ArgumentParser:
         action=BooleanOptionalAction,
         help="displays the task as well",
     )
+    parser.add_argument(
+        "-c",
+        "--cached",
+        default=True,
+        action=BooleanOptionalAction,
+        help="uses cached configuration, only available for some of them, "
+        "mostly for unit test purposes",
+    )
+    parser.add_argument(
+        "--mop",
+        metavar="KEY=VALUE",
+        nargs="*",
+        help="Additional model options, use to change some parameters of the model, "
+        "example: --mop attn_implementation=eager",
+        action=_ParseDict,
+    )
     return parser
 
 
@@ -222,7 +238,11 @@ def _cmd_config(argv: List[Any]):
 
     parser = get_parser_config()
     args = parser.parse_args(argv[1:])
-    print(get_pretrained_config(args.mid))
+    conf = get_pretrained_config(args.mid, **(args.mop or {}))
+    print(conf)
+    for k, v in sorted(conf.__dict__.items()):
+        if "_implementation" in k:
+            print(f"config.{k}={v!r}")
     if args.task:
         print("------")
         print(f"task: {task_from_id(args.mid)}")
@@ -238,6 +258,19 @@ class _ParseDict(argparse.Action):
                 key = split_items[0].strip()  # we remove blanks around keys, as is logical
                 value = split_items[1]
 
+                if value in ("True", "true", "False", "false"):
+                    d[key] = bool(value)
+                    continue
+                try:
+                    d[key] = int(value)
+                    continue
+                except (TypeError, ValueError):
+                    pass
+                try:
+                    d[key] = float(value)
+                    continue
+                except (TypeError, ValueError):
+                    pass
                 d[key] = value
 
         setattr(namespace, self.dest, d)
@@ -321,6 +354,14 @@ def get_parser_validate() -> ArgumentParser:
         "inputs use to export, example: --iop cls_cache=SlidingWindowCache",
         action=_ParseDict,
     )
+    parser.add_argument(
+        "--mop",
+        metavar="KEY=VALUE",
+        nargs="*",
+        help="Additional model options, use to change some parameters of the model, "
+        "example: --mop attn_implementation=eager",
+        action=_ParseDict,
+    )
     return parser
 
 
@@ -371,6 +412,7 @@ def _cmd_validate(argv: List[Any]):
             drop_inputs=None if not args.drop else args.drop.split(","),
             ortfusiontype=args.ortfusiontype,
             input_options=args.iop,
+            model_options=args.mop,
         )
         print("")
         print("-- summary --")

onnx_diagnostic/export/dynamic_shapes.py

@@ -363,16 +363,20 @@ class CoupleInputsDynamicShapes:
         )
         if flatten_unflatten:
             flatunflat = flatten_unflatten_for_dynamic_shapes(inputs)
-            return cls._generic_walker_step(
+            res = cls._generic_walker_step(
                 processor, flatunflat, ds, flatten_unflatten=flatten_unflatten
             )
-        flat, _spec = torch.utils._pytree.tree_flatten(inputs)
+            # Should we restore the original class?
+            return res
+        flat, spec = torch.utils._pytree.tree_flatten(inputs)
         if all(isinstance(t, torch.Tensor) for t in flat):
             # We need to flatten dynamic shapes as well
             ds = flatten_dynamic_shapes(ds)
-        return cls._generic_walker_step(
+        res = cls._generic_walker_step(
             processor, flat, ds, flatten_unflatten=flatten_unflatten
         )
+        # Then we restore the original class.
+        return torch.utils._pytree.tree_unflatten(res, spec)
 
     class ChangeDimensionProcessor:
         def __init__(self, desired_values):

onnx_diagnostic/ext_test_case.py

@@ -461,7 +461,7 @@ def requires_sklearn(version: str, msg: str = "") -> Callable:
     return lambda x: x
 
 
-def requires_experimental(version: str = "", msg: str = "") -> Callable:
+def requires_experimental(version: str = "0.0.0", msg: str = "") -> Callable:
     """Skips a unit test if :epkg:`experimental-experiment` is not recent enough."""
     import packaging.version as pv
 

onnx_diagnostic/helpers/cache_helper.py

@@ -155,6 +155,7 @@ def make_mamba_cache(
     key_value_pairs: List[Tuple[torch.Tensor, torch.Tensor]],
 ) -> transformers.cache_utils.MambaCache:
     "Creates a :class:`transformers.cache_utils.MambaCache`."
+    dtype = key_value_pairs[0][0].dtype
 
     class _config:
         def __init__(self):
@@ -162,14 +163,23 @@ def make_mamba_cache(
             self.conv_kernel = key_value_pairs[0][0].shape[-1]
             self.state_size = key_value_pairs[0][1].shape[-1]
             self.num_hidden_layers = len(key_value_pairs)
-            self.dtype = key_value_pairs[0][0].dtype
+            self.dtype = dtype
 
     cache = transformers.cache_utils.MambaCache(
         _config(),
         max_batch_size=key_value_pairs[0][0].shape[0],
         device=key_value_pairs[0][0].device,
+        dtype=dtype,
     )
     for i in range(len(key_value_pairs)):
+        assert cache.conv_states[i].dtype == dtype, (
+            f"Type mismatch for cache.conv_states[{i}].dtype="
+            f"{cache.conv_states[i].dtype} != {dtype}"
+        )
+        assert cache.ssm_states[i].dtype == dtype, (
+            f"Type mismatch for cache.ssm_states[{i}].dtype="
+            f"{cache.ssm_states[i].dtype} != {dtype}"
+        )
         assert cache.conv_states[i].shape == key_value_pairs[i][0].shape, (
             f"Shape mismatch, expected {cache.conv_states[i].shape}, "
             f"got {key_value_pairs[i][0].shape}"

onnx_diagnostic/helpers/config_helper.py

@@ -28,13 +28,18 @@ def check_hasattr(config: Any, *args: Union[str, Tuple[Any, ...]]):
 def update_config(config: Any, mkwargs: Dict[str, Any]):
     """Updates a configuration with different values."""
     for k, v in mkwargs.items():
+        if k == "attn_implementation":
+            config._attn_implementation = v
+            if getattr(config, "_attn_implementation_autoset", False):
+                config._attn_implementation_autoset = False
+            continue
         if isinstance(v, dict):
             assert hasattr(
                 config, k
             ), f"missing attribute {k!r} in config={config}, cannot update it with {v}"
             update_config(getattr(config, k), v)
-        else:
-            setattr(config, k, v)
+            continue
+        setattr(config, k, v)
 
 
 def _pick(config, *atts):

onnx_diagnostic/helpers/helper.py

@@ -666,6 +666,15 @@ def string_type(
             print(f"[string_type] CACHE4:{type(obj)}")
         return f"{obj.__class__.__name__}(...)"
 
+    if obj.__class__.__name__.endswith("Config"):
+        import transformers.configuration_utils as tcu
+
+        if isinstance(obj, tcu.PretrainedConfig):
+            if verbose:
+                print(f"[string_type] CONFIG:{type(obj)}")
+            s = str(obj.to_diff_dict()).replace("\n", "").replace(" ", "")
+            return f"{obj.__class__.__name__}(**{s})"
+
     if verbose:
         print(f"[string_type] END:{type(obj)}")
     raise AssertionError(f"Unsupported type {type(obj).__name__!r} - {type(obj)}")
@@ -1395,6 +1404,28 @@ def max_diff(
             f"level={level}"
         )
 
+    if expected.__class__.__name__ == "SlidingWindowCache":
+        if got.__class__.__name__ == "SlidingWindowCache":
+            if verbose >= 6:
+                print(f"[max_diff] DynamicCache: {string_type(expected)} ? {string_type(got)}")
+            return max_diff(
+                [expected.key_cache, expected.value_cache],
+                [got.key_cache, got.value_cache],
+                verbose=verbose,
+            )
+        if isinstance(got, tuple) and len(got) == 2:
+            return max_diff(
+                [expected.key_cache, expected.value_cache],
+                [got[0], got[1]],
+                verbose=verbose,
+            )
+        raise AssertionError(
+            f"SlidingWindowCache not fully implemented with classes "
+            f"{expected.__class__.__name__!r} and {got.__class__.__name__!r}, "
+            f"and expected={string_type(expected)}, got={string_type(got)},\n"
+            f"level={level}"
+        )
+
     if expected.__class__.__name__ == "EncoderDecoderCache":
         if got.__class__.__name__ == "EncoderDecoderCache":
             if verbose >= 6:

onnx_diagnostic/helpers/torch_test_helper.py

@@ -8,6 +8,7 @@ from .cache_helper import (
     make_dynamic_cache,
     make_encoder_decoder_cache,
     make_sliding_window_cache,
+    make_mamba_cache,
 )
 
 
@@ -346,6 +347,8 @@ def torch_deepcopy(value: Any) -> Any:
     """
     Makes a deepcopy.
     """
+    if value is None:
+        return None
     if isinstance(value, (int, float, str)):
         return value
     if isinstance(value, tuple):
@@ -376,6 +379,9 @@ def torch_deepcopy(value: Any) -> Any:
             torch_deepcopy(value.self_attention_cache),
             torch_deepcopy(value.cross_attention_cache),
         )
+    if value.__class__.__name__ == "MambaCache":
+        return make_mamba_cache(list(zip(value.conv_states, value.ssm_states)))
+
     if value.__class__ in torch.utils._pytree.SUPPORTED_NODES:
         args, spec = torch.utils._pytree.tree_flatten(value)
         new_args = torch_deepcopy(args)

onnx_diagnostic/tasks/__init__.py

@@ -1,9 +1,11 @@
 from typing import Any, Callable, Dict, List, Tuple
 from . import (
     automatic_speech_recognition,
+    feature_extraction,
     fill_mask,
     image_classification,
     image_text_to_text,
+    mixture_of_expert,
     sentence_similarity,
     text_classification,
     text_generation,
@@ -13,9 +15,11 @@ from . import (
 
 __TASKS__ = [
     automatic_speech_recognition,
+    feature_extraction,
     fill_mask,
     image_classification,
     image_text_to_text,
+    mixture_of_expert,
     sentence_similarity,
     text_classification,
     text_generation,
@@ -33,7 +37,7 @@ def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
     """Reduces a model size."""
     tasks = {mod.__TASK__: mod.reduce_model_config for mod in __TASKS__}
     assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
-    return tasks[task](config, task)
+    return tasks[task](config)
 
 
 def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
@@ -45,4 +49,4 @@ def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callabl
     """
     tasks = {mod.__TASK__: mod.random_input_kwargs for mod in __TASKS__}
     assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
-    return tasks[task](config, task)
+    return tasks[task](config)

onnx_diagnostic/tasks/automatic_speech_recognition.py

@@ -7,7 +7,7 @@ from ..helpers.config_helper import update_config, check_hasattr
 __TASK__ = "automatic-speech-recognition"
 
 
-def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_decoder_layers"):
@@ -33,10 +33,11 @@ def get_inputs(
     head_dim: int,
     batch_size: int = 2,
     sequence_length: int = 30,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
-    Generates inputs for task ``text2text-generation``.
+    Generates inputs for task ``automatic-speech-recognition``.
     Example:
 
     ::
@@ -126,10 +127,27 @@ def get_inputs(
         # encoder_last_hidden_state=torch.randn(batch_size, sequence_length2, encoder_dim),
         # encoder_outputs=torch.randn(batch_size, sequence_length2, encoder_dim),
     )
-    return dict(inputs=inputs, dynamic_shapes=shapes)
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            dummy_max_token_id=dummy_max_token_id,
+            max_source_positions=max_source_positions,
+            d_model=d_model,
+            num_hidden_layers=num_hidden_layers,
+            encoder_attention_heads=encoder_attention_heads,
+            encoder_layers=encoder_layers,
+            decoder_layers=decoder_layers,
+            head_dim=head_dim,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + 1,
+            **kwargs,
+        )["inputs"]
+    return res
 
 
-def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     """
     Inputs kwargs.
 

onnx_diagnostic/tasks/feature_extraction.py

@@ -0,0 +1,76 @@
+from typing import Any, Callable, Dict, Optional, Tuple
+import torch
+from ..helpers.config_helper import update_config, check_hasattr
+
+__TASK__ = "feature-extraction"
+
+
+def reduce_model_config(config: Any) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(config, "num_attention_heads", "num_hidden_layers")
+    kwargs = dict(
+        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_attention_heads=min(config.num_attention_heads, 4),
+    )
+    update_config(config, kwargs)
+    return kwargs
+
+
+def get_inputs(
+    model: torch.nn.Module,
+    config: Optional[Any],
+    batch_size: int,
+    sequence_length: int,
+    dummy_max_token_id: int,
+    add_second_input: bool = False,
+    **kwargs,  # unused
+):
+    """
+    Generates inputs for task ``feature-extraction``.
+    Example:
+
+    ::
+
+        input_ids:T7s1x13[101,72654:A16789.23076923077],
+        token_type_ids:T7s1x13[0,0:A0.0],
+        attention_mask:T7s1x13[1,1:A1.0])
+    """
+    batch = torch.export.Dim("batch", min=1, max=1024)
+    seq_length = "sequence_length"
+    shapes = {
+        "input_ids": {0: batch, 1: seq_length},
+        "attention_mask": {0: batch, 1: seq_length},
+    }
+    inputs = dict(
+        input_ids=torch.randint(0, dummy_max_token_id, (batch_size, sequence_length)).to(
+            torch.int64
+        ),
+        attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64),
+    )
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + 1,
+            dummy_max_token_id=dummy_max_token_id,
+            **kwargs,
+        )["inputs"]
+    return res
+
+
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
+    """
+    Inputs kwargs.
+
+    If the configuration is None, the function selects typical dimensions.
+    """
+    if config is not None:
+        check_hasattr(config, "vocab_size")
+    kwargs = dict(
+        batch_size=2,
+        sequence_length=30,
+        dummy_max_token_id=31999 if config is None else (config.vocab_size - 1),
+    )
+    return kwargs, get_inputs

onnx_diagnostic/tasks/fill_mask.py

@@ -5,7 +5,7 @@ from ..helpers.config_helper import update_config, check_hasattr
 __TASK__ = "fill-mask"
 
 
-def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
@@ -22,6 +22,7 @@ def get_inputs(
     batch_size: int,
     sequence_length: int,
     dummy_max_token_id: int,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -48,10 +49,20 @@ def get_inputs(
         token_type_ids=torch.zeros((batch_size, sequence_length)).to(torch.int64),
         attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64),
     )
-    return dict(inputs=inputs, dynamic_shapes=shapes)
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + 1,
+            dummy_max_token_id=dummy_max_token_id,
+            **kwargs,
+        )["inputs"]
+    return res
 
 
-def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     """
     Inputs kwargs.
 

onnx_diagnostic/tasks/image_classification.py

@@ -5,7 +5,7 @@ from ..helpers.config_helper import update_config, check_hasattr
 __TASK__ = "image-classification"
 
 
-def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
@@ -27,6 +27,7 @@ def get_inputs(
     input_channels: int,
     batch_size: int = 2,
     dynamic_rope: bool = False,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -59,10 +60,22 @@ def get_inputs(
             -1, 1
         ),
     )
-    return dict(inputs=inputs, dynamic_shapes=shapes)
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            input_width=input_width + 1,
+            input_height=input_height + 1,
+            input_channels=input_channels,
+            batch_size=batch_size + 1,
+            dynamic_rope=dynamic_rope,
+            **kwargs,
+        )["inputs"]
+    return res
 
 
-def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     """
     Inputs kwargs.
 

onnx_diagnostic/tasks/image_text_to_text.py

@@ -6,7 +6,7 @@ from ..helpers.config_helper import update_config, check_hasattr, _pick
 __TASK__ = "image-text-to-text"
 
 
-def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     kwargs: Dict[str, Any] = {}
     if hasattr(config, "num_hidden_layers"):
@@ -32,10 +32,11 @@ def get_inputs(
     sequence_length2: int = 3,
     n_images: int = 2,
     dynamic_rope: bool = False,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
-    Generates input for task ``text-generation``.
+    Generates input for task ``image-text-to-text``.
 
     :param model: model to get the missing information
     :param config: configuration used to generate the model
@@ -99,10 +100,29 @@ def get_inputs(
             torch.int64
         ),
     )
-    return dict(inputs=inputs, dynamic_shapes=shapes)
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            dummy_max_token_id=dummy_max_token_id,
+            num_key_value_heads=num_key_value_heads,
+            num_hidden_layers=num_hidden_layers,
+            head_dim=head_dim,
+            width=width,
+            height=height,
+            num_channels=num_channels,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + 1,
+            sequence_length2=sequence_length2 + 1,
+            n_images=n_images + 1,
+            dynamic_rope=dynamic_rope,
+            **kwargs,
+        )["inputs"]
+    return res
 
 
-def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     """
     Inputs kwargs.
 

onnx_diagnostic/tasks/mixture_of_expert.py

@@ -0,0 +1,76 @@
+from typing import Any, Callable, Dict, Optional, Tuple
+import torch
+
+# from ..helpers.cache_helper import make_dynamic_cache
+from ..helpers.config_helper import update_config  # , check_hasattr, _pick
+
+__TASK__ = "MoE"
+
+
+def reduce_model_config(config: Any) -> Dict[str, Any]:
+    """Reduces a model size."""
+    kwargs: Dict[str, Any] = {}
+    if hasattr(config, "num_hidden_layers"):
+        config.num_hidden_layers = min(config.num_hidden_layers, 2)
+    if hasattr(config, "vision_config") and hasattr(config.vision_config, "num_hidden_layers"):
+        config.vision_config.num_hidden_layers = min(config.vision_config.num_hidden_layers, 2)
+    if hasattr(config, "audio_processor") and hasattr(
+        config.audio_processor, "num_hidden_layers"
+    ):
+        config.audio_processor.num_hidden_layers = min(
+            config.audio_processor.num_hidden_layers, 2
+        )
+    if hasattr(config, "audio_processor") and hasattr(config.audio_processor, "attention_dim"):
+        config.audio_processor.attention_dim = min(config.audio_processor.attention_dim, 2)
+    update_config(config, kwargs)
+    return kwargs
+
+
+def get_inputs(
+    model: torch.nn.Module,
+    config: Optional[Any],
+    dummy_max_token_id: int,
+    num_key_value_heads: int,
+    num_hidden_layers: int,
+    head_dim: int,
+    width: int,
+    height: int,
+    num_channels: int,
+    batch_size: int = 2,
+    sequence_length: int = 30,
+    sequence_length2: int = 3,
+    n_images: int = 2,
+    dynamic_rope: bool = False,
+    add_second_input: bool = False,
+    **kwargs,  # unused
+):
+    """
+    Generates input for task ``MoE``.
+
+    :param model: model to get the missing information
+    :param config: configuration used to generate the model
+    :param head_dim: last dimension of the cache
+    :param dummy_max_token_id: dummy max token id
+    :param batch_size: batch size
+    :param sequence_length: sequence length
+    :param sequence_length2: new sequence length
+    :param n_images: number of images
+    :param width: width of the image
+    :param height: height of the image
+    :param num_channels: number of channels
+    :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
+    :return: dictionary
+    """
+    assert not add_second_input, "add_second_input=True not yet implemented"
+    raise NotImplementedError(f"get_inputs not yet implemented for task {__TASK__!r}.")
+
+
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
+    """
+    Inputs kwargs.
+
+    If the configuration is None, the function selects typical dimensions.
+    """
+    raise NotImplementedError(
+        f"random_input_kwargs not yet implemented for task {__TASK__!r}."
+    )

onnx_diagnostic/tasks/sentence_similarity.py

@@ -5,7 +5,7 @@ from ..helpers.config_helper import update_config, check_hasattr
 __TASK__ = "sentence-similarity"
 
 
-def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     check_hasattr(config, "num_attention_heads", "num_hidden_layers")
     kwargs = dict(
@@ -22,6 +22,7 @@ def get_inputs(
     batch_size: int,
     sequence_length: int,
     dummy_max_token_id: int,
+    add_second_input: bool = False,
     **kwargs,  # unused
 ):
     """
@@ -48,10 +49,20 @@ def get_inputs(
         token_type_ids=torch.zeros((batch_size, sequence_length)).to(torch.int64),
         attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64),
     )
-    return dict(inputs=inputs, dynamic_shapes=shapes)
+    res = dict(inputs=inputs, dynamic_shapes=shapes)
+    if add_second_input:
+        res["inputs2"] = get_inputs(
+            model=model,
+            config=config,
+            batch_size=batch_size + 1,
+            sequence_length=sequence_length + 1,
+            dummy_max_token_id=dummy_max_token_id,
+            **kwargs,
+        )["inputs"]
+    return res
 
 
-def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     """
     Inputs kwargs.