onnx-diagnostic 0.7.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

onnx_diagnostic/helpers/log_helper.py

@@ -21,6 +21,63 @@ BUCKET_SCALES_VALUES = np.array(
 BUCKET_SCALES = BUCKET_SCALES_VALUES / 100 + 1
 
 
+def filter_data(
+    df: pandas.DataFrame,
+    filter_in: Optional[str] = None,
+    filter_out: Optional[str] = None,
+    verbose: int = 0,
+) -> pandas.DataFrame:
+    """
+    Argument `filter` follows the syntax
+    ``<column1>:<fmt1>//<column2>:<fmt2>``.
+
+    The format is the following:
+
+    * a value or a set of values separated by ``;``
+    """
+    if not filter_in and not filter_out:
+        return df
+
+    def _f(fmt):
+        cond = {}
+        if isinstance(fmt, str):
+            cols = fmt.split("//")
+            for c in cols:
+                assert ":" in c, f"Unexpected value {c!r} in fmt={fmt!r}"
+                spl = c.split(":")
+                assert len(spl) == 2, f"Unexpected value {c!r} in fmt={fmt!r}"
+                name, fil = spl
+                cond[name] = set(fil.split(";"))
+        return cond
+
+    if filter_in:
+        cond = _f(filter_in)
+        assert isinstance(cond, dict), f"Unexpected type {type(cond)} for fmt={filter_in!r}"
+        for k, v in cond.items():
+            if k not in df.columns:
+                continue
+            if verbose:
+                print(
+                    f"[_filter_data] filter in column {k!r}, "
+                    f"values {v!r} among {set(df[k].astype(str))}"
+                )
+            df = df[df[k].astype(str).isin(v)]
+
+    if filter_out:
+        cond = _f(filter_out)
+        assert isinstance(cond, dict), f"Unexpected type {type(cond)} for fmt={filter_out!r}"
+        for k, v in cond.items():
+            if k not in df.columns:
+                continue
+            if verbose:
+                print(
+                    f"[_filter_data] filter out column {k!r}, "
+                    f"values {v!r} among {set(df[k].astype(str))}"
+                )
+            df = df[~df[k].astype(str).isin(v)]
+    return df
+
+
 def enumerate_csv_files(
     data: Union[
         pandas.DataFrame, List[Union[str, Tuple[str, str]]], str, Tuple[str, str, str, str]
@@ -118,7 +175,8 @@ def open_dataframe(
     data: Union[str, Tuple[str, str, str, str], pandas.DataFrame],
 ) -> pandas.DataFrame:
     """
-    Opens a filename.
+    Opens a filename defined by function
+    :func:`onnx_diagnostic.helpers.log_helper.enumerate_csv_files`.
 
     :param data: a dataframe, a filename, a tuple indicating the file is coming
         from a zip file
@@ -259,7 +317,7 @@ def apply_excel_style(
         co: Dict[int, int] = {}
         sizes: Dict[int, int] = {}
         cols = set()
-        for i in range(1, n_rows):
+        for i in range(1, n_rows + 1):
             for j, cell in enumerate(sheet[i]):
                 if j > n_cols:
                     break
@@ -277,7 +335,7 @@ def apply_excel_style(
                 c = get_column_letter(k)
                 sheet.column_dimensions[c].width = 15
 
-        for i in range(1, n_rows):
+        for i in range(1, n_rows + 1):
             for j, cell in enumerate(sheet[i]):
                 if j > n_cols:
                     break
@@ -333,18 +391,85 @@ def apply_excel_style(
 class CubePlot:
     """
     Creates a plot.
+
+    :param df: dataframe
+    :param kind: kind of graph to plot, bar, barh, line
+    :param split: draw a graph per line in the dataframe
+    :param timeseries: this assumes the time is one level of the columns,
+        this argument indices the level name
     """
 
+    KINDS = {"bar", "barh", "line"}
+
+    @classmethod
+    def group_columns(
+        cls, columns: List[str], sep: str = "/", depth: int = 2
+    ) -> List[List[str]]:
+        """Groups columns to have nice display."""
+        res: Dict[str, List[str]] = {}
+        for c in columns:
+            p = c.split("/")
+            k = "/".join(p[:depth])
+            if k not in res:
+                res[k] = []
+            res[k].append(c)
+        new_res: Dict[str, List[str]] = {}
+        for k, v in res.items():
+            if len(v) >= 3:
+                new_res[k] = v
+            else:
+                if "0" not in new_res:
+                    new_res["0"] = []
+                new_res["0"].extend(v)
+        groups: List[List[str]] = [sorted(v) for k, v in sorted(new_res.items())]
+        if depth <= 1:
+            return groups
+        new_groups: List[List[str]] = []
+        for v in groups:
+            if len(v) >= 6:
+                new_groups.extend(cls.group_columns(v, depth=1, sep=sep))
+            else:
+                new_groups.append(v)
+        return new_groups
+
     def __init__(
-        self, df: pandas.DataFrame, kind: str = "bar", orientation="col", split: bool = True
+        self,
+        df: pandas.DataFrame,
+        kind: str = "bar",
+        orientation="col",
+        split: bool = True,
+        timeseries: Optional[str] = None,
     ):
+        assert (
+            not timeseries or timeseries in df.columns.names
+        ), f"Level {timeseries!r} is not part of the columns levels {df.columns.names}"
+        assert (
+            kind in self.__class__.KINDS
+        ), f"Unexpected kind={kind!r} not in {self.__class__.KINDS}"
+        assert split, f"split={split} not implemented"
+        assert (
+            not timeseries or orientation == "row"
+        ), f"orientation={orientation!r} must be 'row' for timeseries"
         self.df = df.copy()
         self.kind = kind
         self.orientation = orientation
         self.split = split
+        self.timeseries = timeseries
+
+        if timeseries:
+            if isinstance(self.df.columns, pandas.MultiIndex):
+                index_time = list(self.df.columns.names).index(self.timeseries)
 
-        if isinstance(self.df.columns, pandas.MultiIndex):
-            self.df.columns = ["/".join(map(str, i)) for i in self.df.columns]
+                def _drop(t, i=index_time):
+                    return (*t[:i], *t[i + 1 :])
+
+                self.df.columns = pandas.MultiIndex.from_tuples(
+                    [("/".join(map(str, _drop(i))), i[index_time]) for i in self.df.columns],
+                    names=["metric", timeseries],
+                )
+        else:
+            if isinstance(self.df.columns, pandas.MultiIndex):
+                self.df.columns = ["/".join(map(str, i)) for i in self.df.columns]
         if isinstance(self.df.index, pandas.MultiIndex):
             self.df.index = ["/".join(map(str, i)) for i in self.df.index]
 
@@ -354,85 +479,129 @@ class CubePlot:
 
     def to_images(
         self, verbose: int = 0, merge: bool = True, title_suffix: Optional[str] = None
-    ):
+    ) -> List[bytes]:
         """
         Converts data into plots and images.
+
+        :param verbose: verbosity
+        :param merge: returns all graphs in a single image (True)
+            or an image for every graph (False)
+        :param title_suffix: prefix for the title of every graph
+        :return: list of binary images (format PNG)
         """
-        import matplotlib.pyplot as plt
+        if self.kind in ("barh", "bar"):
+            return self._to_images_bar(verbose=verbose, merge=merge, title_suffix=title_suffix)
+        if self.kind == "line":
+            return self._to_images_line(
+                verbose=verbose, merge=merge, title_suffix=title_suffix
+            )
+        raise AssertionError(f"self.kind={self.kind!r} not implemented")
 
-        df = self.df.T if self.orientation == "row" else self.df
-        imgs = []
+    @classmethod
+    def _make_loop(cls, ensemble, verbose):
         if verbose:
             from tqdm import tqdm
 
-            loop = tqdm(df.columns)
+            loop = tqdm(ensemble)
         else:
-            loop = df.columns
+            loop = ensemble
+        return loop
+
+    def _to_images_bar(
+        self, verbose: int = 0, merge: bool = True, title_suffix: Optional[str] = None
+    ) -> List[bytes]:
+        assert merge, f"merge={merge} not implemented yet"
+        import matplotlib.pyplot as plt
+
+        df = self.df.T if self.orientation == "row" else self.df
         title_suffix = f"\n{title_suffix}" if title_suffix else ""
-        if merge:
-            nn = len(df.columns) // 2
-            nn += nn % 2
-            fig, axs = plt.subplots(nn, 2, figsize=(12, 3 * nn * df.shape[0] / 12))
-            pos = 0
-            for c in loop:
-                ax = axs[pos // 2, pos % 2]
+
+        n_cols = 3
+        nn = df.shape[1] // n_cols
+        nn += int(df.shape[1] % n_cols != 0)
+        fig, axs = plt.subplots(nn, n_cols, figsize=(6 * n_cols, nn * df.shape[0] / 5))
+        pos = 0
+        imgs = []
+        for c in self._make_loop(df.columns, verbose):
+            ax = axs[pos // n_cols, pos % n_cols]
+            (
                 df[c].plot.barh(title=f"{c}{title_suffix}", ax=ax)
-                ax.tick_params(axis="both", which="major", labelsize=8)
-                ax.grid(True)
-                pos += 1  # noqa: SIM113
-            fig.tight_layout()
-            imgdata = io.BytesIO()
-            fig.savefig(imgdata, format="png")
-            imgs.append(imgdata.getvalue())
-            plt.close()
-        else:
-            for c in loop:
-                fig, ax = plt.subplots(1, 1, figsize=(3, 3))
-                df[c].plot.barh(title=c, ax=ax)
-                ax.tick_params(axis="both", which="major", labelsize=8)
-                ax.grid(True)
-                fig.tight_layout()
-                imgdata = io.BytesIO()
-                fig.savefig(imgdata, format="png")
-                imgs.append(imgdata.getvalue())
-                plt.close()
+                if self.kind == "barh"
+                else df[c].plot.bar(title=f"{c}{title_suffix}", ax=ax)
+            )
+            ax.tick_params(axis="both", which="major", labelsize=8)
+            ax.grid(True)
+            pos += 1  # noqa: SIM113
+        fig.tight_layout()
+        imgdata = io.BytesIO()
+        fig.savefig(imgdata, format="png")
+        imgs.append(imgdata.getvalue())
+        plt.close()
         return imgs
 
-    def to_charts(self, writer: pandas.ExcelWriter, sheet, empty_row: int = 1):
-        """
-        Draws plots on a page.
-        The data is copied on this page.
-
-        :param name: sheet name
-        :param writer: writer (from pandas)
-        :param sheet_name: sheet
-        :param graph_index: graph index
-        :return: list of graph
-        """
-        assert self.split, f"Not implemented if split={self.split}"
-        assert self.orientation == "row", f"Not implemented if orientation={self.orientation}"
-        workbook = writer.book
-        labels = list(self.df.columns)
-        sheet.write_row(empty_row, 0, labels)
-
-        charts = []
-        pos = empty_row + 1
-        for i in self.df.index:
-            values = self.df.loc[i, :].tolist()
-            values = [("" if isinstance(v, float) and np.isnan(v) else v) for v in values]
-            sheet.write_row(pos, 0, values)
-            chart = workbook.add_chart({"type": "bar"})
-            chart.add_series(
-                {
-                    "name": i,
-                    "categories": [i, 1, empty_row, len(labels), empty_row],
-                    "values": [i, 1, pos, len(labels), pos],
-                }
-            )
-            chart.set_title({"name": i})
-            charts.append(chart)
-            pos += 1
-        return charts
+    def _to_images_line(
+        self, verbose: int = 0, merge: bool = True, title_suffix: Optional[str] = None
+    ) -> List[bytes]:
+        assert merge, f"merge={merge} not implemented yet"
+        assert (
+            self.orientation == "row"
+        ), f"self.orientation={self.orientation!r} not implemented for this kind of graph."
+
+        def rotate_align(ax, angle=15, align="right"):
+            for label in ax.get_xticklabels():
+                label.set_rotation(angle)
+                label.set_horizontalalignment(align)
+            ax.tick_params(axis="both", which="major", labelsize=8)
+            ax.grid(True)
+            ax.legend()
+            ax.tick_params(labelleft=True)
+            return ax
+
+        import matplotlib.pyplot as plt
+
+        df = self.df.T
+
+        confs = list(df.unstack(self.timeseries).index)
+        groups = self.group_columns(confs)
+        n_cols = len(groups)
+
+        title_suffix = f"\n{title_suffix}" if title_suffix else ""
+        fig, axs = plt.subplots(
+            df.shape[1],
+            n_cols,
+            figsize=(5 * n_cols, max(len(g) for g in groups) * df.shape[1] / 2),
+            sharex=True,
+            sharey="row" if n_cols > 1 else False,
+        )
+        imgs = []
+        row = 0
+        for c in self._make_loop(df.columns, verbose):
+            dfc = df[[c]]
+            dfc = dfc.unstack(self.timeseries).T.droplevel(0)
+            if n_cols == 1:
+                dfc.plot(title=f"{c}{title_suffix}", ax=axs[row], linewidth=3)
+                axs[row].grid(True)
+                rotate_align(axs[row])
+            else:
+                x = list(range(dfc.shape[0]))
+                ticks = list(dfc.index)
+                for ii, group in enumerate(groups):
+                    ddd = dfc.loc[:, group].copy()
+                    axs[row, ii].set_xticks(x)
+                    axs[row, ii].set_xticklabels(ticks)
+                    # This is very slow
+                    # ddd.plot(ax=axs[row, ii],linewidth=3)
+                    for jj in range(ddd.shape[1]):
+                        axs[row, ii].plot(x, ddd.iloc[:, jj], lw=3, label=ddd.columns[jj])
+                    axs[row, ii].set_title(f"{c}{title_suffix}")
+                    rotate_align(axs[row, ii])
+            row += 1  # noqa: SIM113
+        fig.tight_layout()
+        imgdata = io.BytesIO()
+        fig.savefig(imgdata, format="png")
+        imgs.append(imgdata.getvalue())
+        plt.close()
+        return imgs
 
 
 class CubeLogs:
@@ -1084,7 +1253,11 @@ class CubeLogs:
                 df.to_excel(writer, sheet_name=main, freeze_panes=(1, 1))
 
             for name, view in views.items():
+                if view is None:
+                    continue
                 df, tview = self.view(view, return_view_def=True, verbose=max(verbose - 1, 0))
+                if tview is None:
+                    continue
                 memory = df.memory_usage(deep=True).sum()
                 if verbose:
                     print(
@@ -1128,7 +1301,17 @@ class CubeLogs:
                     )
                     f_highlights[name] = tview.f_highlight
                     if tview.plots:
-                        plots.append(CubePlot(df, kind="barh", orientation="row", split=True))
+                        plots.append(
+                            CubePlot(
+                                df,
+                                kind="line",
+                                orientation="row",
+                                split=True,
+                                timeseries=self.time,
+                            )
+                            if self.time in df.columns.names
+                            else CubePlot(df, kind="barh", orientation="row", split=True)
+                        )
             if raw:
                 assert main not in views, f"{main!r} is duplicated in views {sorted(views)}"
                 # Too long.
@@ -1249,9 +1432,11 @@ class CubeLogsPerformance(CubeLogs):
             "n_node_scatter",
             "n_node_function",
             "n_node_initializer",
+            "n_node_initializer_small",
             "n_node_constant",
             "n_node_shape",
             "n_node_expand",
+            "onnx_n_nodes_no_cst",
             "peak_gpu_torch",
             "peak_gpu_nvidia",
             "time_export_unbiased",
@@ -1419,6 +1604,9 @@ class CubeLogsPerformance(CubeLogs):
                 n_node_function=lambda df: gpreserve(
                     df, "onnx_n_functions", gdf(df, "onnx_n_functions")
                 ),
+                n_node_initializer_small=lambda df: gpreserve(
+                    df, "op_onnx_initializer_small", gdf(df, "op_onnx_initializer_small")
+                ),
                 n_node_initializer=lambda df: gpreserve(
                     df, "onnx_n_initializer", gdf(df, "onnx_n_initializer")
                 ),
@@ -1437,6 +1625,10 @@ class CubeLogsPerformance(CubeLogs):
             ), f"Unexpected formula={formula!r}, should be in {sorted(lambdas)}"
             return lambdas[formula]
 
+        if formula == "onnx_n_nodes_no_cst":
+            return lambda df: gdf(df, "onnx_n_nodes", 0) - gdf(
+                df, "op_onnx__Constant", 0
+            ).fillna(0)
         if formula == "peak_gpu_torch":
             return lambda df: gdf(df, "mema_gpu_5_after_export") - gdf(df, "mema_gpu_4_reset")
         if formula == "peak_gpu_nvidia":
@@ -1466,10 +1658,12 @@ class CubeLogsPerformance(CubeLogs):
 
     def view(
         self,
-        view_def: Union[str, CubeViewDef],
+        view_def: Optional[Union[str, CubeViewDef]],
         return_view_def: bool = False,
         verbose: int = 0,
-    ) -> Union[pandas.DataFrame, Tuple[pandas.DataFrame, CubeViewDef]]:
+    ) -> Union[
+        Optional[pandas.DataFrame], Tuple[Optional[pandas.DataFrame], Optional[CubeViewDef]]
+    ]:
         """
         Returns a dataframe, a pivot view.
 
@@ -1478,18 +1672,22 @@ class CubeLogsPerformance(CubeLogs):
         :param view_def: view definition or a string
         :param return_view_def: returns the view definition as well
         :param verbose: verbosity level
-        :return: dataframe
+        :return: dataframe or a couple (dataframe, view definition),
+            both of them can be one if view_def cannot be interpreted
         """
+        assert view_def is not None, "view_def is None, this is not allowed."
         if isinstance(view_def, str):
             view_def = self.make_view_def(view_def)
+            if view_def is None:
+                return (None, None) if return_view_def else None
         return super().view(view_def, return_view_def=return_view_def, verbose=verbose)
 
-    def make_view_def(self, name: str) -> CubeViewDef:
+    def make_view_def(self, name: str) -> Optional[CubeViewDef]:
         """
         Returns a view definition.
 
         :param name: name of the view
-        :return: a CubeViewDef
+        :return: a CubeViewDef or None if name does not make sense
 
         Available views:
 
@@ -1588,6 +1786,8 @@ class CubeLogsPerformance(CubeLogs):
                         "onnx_weight_size_torch",
                         "onnx_weight_size_proto",
                         "onnx_n_nodes",
+                        "onnx_n_nodes_no_cst",
+                        "op_onnx__Constant",
                         "peak_gpu_torch",
                         "peak_gpu_nvidia",
                     ],
@@ -1617,6 +1817,7 @@ class CubeLogsPerformance(CubeLogs):
                         "onnx_weight_size_torch",
                         "onnx_weight_size_proto",
                         "onnx_n_nodes",
+                        "onnx_n_nodes_no_cst",
                         "peak_gpu_torch",
                         "peak_gpu_nvidia",
                     ],
@@ -1701,12 +1902,22 @@ class CubeLogsPerformance(CubeLogs):
                 f_highlight=f_bucket,
                 order=order,
             ),
-            "cmd": lambda: CubeViewDef(
+            "onnx": lambda: CubeViewDef(
                 key_index=index_cols,
-                values=self._filter_column(["CMD"], self.values),
+                values=self._filter_column(
+                    [
+                        "onnx_filesize",
+                        "onnx_n_nodes",
+                        "onnx_n_nodes_no_cst",
+                        "onnx_weight_size_proto",
+                        "onnx_weight_size_torch",
+                        "op_onnx_initializer_small",
+                    ],
+                    self.values,
+                ),
                 ignore_unique=True,
                 keep_columns_in_index=["suite"],
-                name="cmd",
+                name="onnx",
                 order=order,
             ),
             "raw-short": lambda: CubeViewDef(
@@ -1718,11 +1929,25 @@ class CubeLogsPerformance(CubeLogs):
                 no_index=True,
             ),
         }
-        assert name in implemented_views, (
+
+        cmd_col = self._filter_column(["CMD"], self.values, can_be_empty=True)
+        if cmd_col:
+            implemented_views["cmd"] = lambda: CubeViewDef(
+                key_index=index_cols,
+                values=cmd_col,
+                ignore_unique=True,
+                keep_columns_in_index=["suite"],
+                name="cmd",
+                order=order,
+            )
+
+        assert name in implemented_views or name in {"cmd"}, (
             f"Unknown view {name!r}, expected a name in {sorted(implemented_views)},"
             f"\n--\nkeys={pprint.pformat(sorted(self.keys_time))}, "
             f"\n--\nvalues={pprint.pformat(sorted(self.values))}"
         )
+        if name not in implemented_views:
+            return None
         return implemented_views[name]()
 
     def post_load_process_piece(

onnx_diagnostic/helpers/rt_helper.py

@@ -112,4 +112,14 @@ def make_feeds(
 
     if copy:
         flat = [t.copy() if hasattr(t, "copy") else t.clone() for t in flat]
-    return dict(zip(names, flat))
+    # bool, int, float, onnxruntime does not support float, bool, int
+    new_flat = []
+    for i in flat:
+        if isinstance(i, bool):
+            i = np.array(i, dtype=np.bool_)
+        elif isinstance(i, int):
+            i = np.array(i, dtype=np.int64)
+        elif isinstance(i, float):
+            i = np.array(i, dtype=np.float32)
+        new_flat.append(i)
+    return dict(zip(names, new_flat))

onnx_diagnostic/helpers/torch_helper.py

@@ -717,7 +717,7 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
         return tuple(to_any(t, to_value) for t in value)
     if isinstance(value, set):
         return {to_any(t, to_value) for t in value}
-    if isinstance(value, dict):
+    if type(value) is dict:
         return {k: to_any(t, to_value) for k, t in value.items()}
     if value.__class__.__name__ == "DynamicCache":
         return make_dynamic_cache(
@@ -735,7 +735,8 @@ def to_any(value: Any, to_value: Union[torch.dtype, torch.device, str]) -> Any:
                     [t.to(to_value) for t in value.key_cache],
                     [t.to(to_value) for t in value.value_cache],
                 )
-            )
+            ),
+            max_cache_len=value.max_cache_len,
         )
     if value.__class__.__name__ == "EncoderDecoderCache":
         return make_encoder_decoder_cache(
@@ -784,7 +785,10 @@ def torch_deepcopy(value: Any) -> Any:
             torch_deepcopy(list(zip(value.key_cache, value.value_cache)))
         )
     if value.__class__.__name__ == "StaticCache":
-        return make_static_cache(torch_deepcopy(list(zip(value.key_cache, value.value_cache))))
+        return make_static_cache(
+            torch_deepcopy(list(zip(value.key_cache, value.value_cache))),
+            max_cache_len=value.max_cache_len,
+        )
     if value.__class__.__name__ == "SlidingWindowCache":
         return make_sliding_window_cache(
             torch_deepcopy(list(zip(value.key_cache, value.value_cache)))

onnx_diagnostic/tasks/__init__.py

@@ -11,6 +11,7 @@ from . import (
     summarization,
     text_classification,
     text_generation,
+    text_to_image,
     text2text_generation,
     zero_shot_image_classification,
 )
@@ -27,6 +28,7 @@ __TASKS__ = [
     summarization,
     text_classification,
     text_generation,
+    text_to_image,
     text2text_generation,
     zero_shot_image_classification,
 ]

onnx_diagnostic/tasks/text_generation.py

@@ -109,7 +109,7 @@ def get_inputs(
         sequence_length2 = seq_length_multiple
 
         shapes = {
-            "input_ids": {0: batch, 1: torch.export.Dim.DYNAMIC},
+            "input_ids": {0: batch, 1: "sequence_length"},
             "attention_mask": {
                 0: batch,
                 1: "cache+seq",  # cache_length + seq_length
@@ -176,8 +176,10 @@ def get_inputs(
                 "attention_mask": {0: batch, 2: "seq"},
                 "cache_position": {0: "seq"},
                 "past_key_values": [
-                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
-                    [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                    # [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                    # [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],
+                    [{0: batch} for _ in range(num_hidden_layers)],
+                    [{0: batch} for _ in range(num_hidden_layers)],
                 ],
             }
             inputs = dict(
@@ -188,18 +190,25 @@ def get_inputs(
                     (batch_size, num_key_value_heads, sequence_length2, head_dim)
                 ).to(torch.bool),
                 cache_position=torch.arange(sequence_length2).to(torch.int64),
-                past_key_values=make_cache(
+                past_key_values=make_static_cache(
                     [
                         (
                             torch.randn(
-                                batch_size, num_key_value_heads, sequence_length, head_dim
+                                batch_size,
+                                num_key_value_heads,
+                                sequence_length + sequence_length2,
+                                head_dim,
                             ),
                             torch.randn(
-                                batch_size, num_key_value_heads, sequence_length, head_dim
+                                batch_size,
+                                num_key_value_heads,
+                                sequence_length + sequence_length2,
+                                head_dim,
                             ),
                         )
                         for i in range(num_hidden_layers)
-                    ]
+                    ],
+                    max_cache_len=max(sequence_length + sequence_length2, head_dim),
                 ),
             )
         else:
@@ -230,7 +239,7 @@ def get_inputs(
                 position_ids=torch.arange(sequence_length, sequence_length + sequence_length2)
                 .to(torch.int64)
                 .expand((batch_size, -1)),
-                past_key_values=make_cache(
+                past_key_values=make_cache(  # type: ignore[operator]
                     [
                         (
                             torch.randn(