ut

Author: xadupre

_unittests/ut_export/test_jit.py

@@ -1,3 +1,4 @@
+import inspect
 import unittest
 import torch
 from onnx_diagnostic.ext_test_case import (
@@ -15,6 +16,9 @@
     to_onnx = None
 
 
+has_scan_reverse = "reverse" in set(inspect.signature(torch.ops.higher_order.scan).parameters)
+
+
 @torch.jit.script_if_tracing
 def dummy_loop(padded: torch.Tensor, pos: torch.Tensor):
     copy = torch.zeros(padded.shape)
@@ -36,12 +40,12 @@ def pad_row(padded, p):
         row[: p.item()] = padded[: p.item()]
         return (row,)
 
-    return torch.ops.higher_order.scan(
-        pad_row,
-        [],
-        [padded, pos],
-        [],
-    )
+    if has_scan_reverse:
+        # torch==2.6
+        return torch.ops.higher_order.scan(
+            pad_row, [], [padded, pos], additional_inputs=[], reverse=False, dim=0
+        )
+    return torch.ops.higher_order.scan(pad_row, [], [padded, pos], [])
 
 
 def select_when_exporting(f, f_scan):