Remove if branches in arange | fix(torchlib) (#1097)

Replace if with where for performance.
microsoft · Oct 23, 2023 · b0a0b87 · b0a0b87
1 parent 778e8e8
commit b0a0b87
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -591,11 +591,8 @@ def _adjust_args_for_arange_int_dtype(
     end = op.Cast(end, to=FLOAT.dtype)
     step = op.Cast(step, to=FLOAT.dtype)
 
-    if start < zero:
-        start = op.Ceil(start)
-
-    if step < zero:
-        start = op.Floor(start)
+    start = op.Where(op.Less(start, zero), op.Ceil(start), start)
+    start = op.Where(op.Less(step, zero), op.Floor(start), start)
 
     return (start, end, step)