Fix eager evaluation of the Rank op | fix(torchlib) (#1303)

I mistakenly copied the logic for `IsScalar` to `Rank`. `Rank` needs its own logic.
microsoft · Mar 14, 2024 · 7f53395 · 7f53395
1 parent d681dbd
commit 7f53395
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/onnxscript/function_libs/torch_lib/graph_building.py b/onnxscript/function_libs/torch_lib/graph_building.py
@@ -366,7 +366,11 @@ def eval_function(  # type: ignore[override]
                         # Fall to call add_function_call
                         pass
                 elif isinstance(args[0], Sequence):
-                    return False
+                    if all(isinstance(arg, (int, float)) for arg in args[0]):
+                        return 1
+                    else:
+                        # Fall to call add_function_call
+                        pass
                 else:
                     # Python constants are scalars
                     return 0