Update base for Update on "Mark some functions as traceable | feat(torchlib)"

As an effort described in #1095, this PR marks functions with if branches as traceable.

[ghstack-poisoned]

Author: justinchuby

onnxscript/function_libs/torch_lib/ops/core.py

@@ -365,7 +365,9 @@ def aten_all_dims(self: TTensor, dim: Sequence[int] = (), keepdim: bool = False)
     if not dim:
         return aten_all_dims_no_dim(self, keepdim)
     for d in dim:
-        self = aten_all_dim(self, d, keepdim)
+        self = aten_all_dim(self, d, keepdim=True)
+    if not keepdim:
+        self = op.Squeeze(self, list(dim))
     return self
 
 
@@ -488,7 +490,9 @@ def aten_any_dims(self: TTensor, dim: Sequence[int] = (), keepdim: bool = False)
     if not dim:
         return aten_any_dims_no_dim(self, keepdim)
     for d in dim:
-        self = aten_any_dim(self, d, keepdim)
+        self = aten_any_dim(self, d, keepdim=True)
+    if not keepdim:
+        self = op.Squeeze(self, list(dim))
     return self
 
 

onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -483,9 +483,14 @@ def _where_input_wrangler(
         reason="PyTorch does not implement _softmax for float16 on CPU",
         dtypes=(torch.float16,),
     ),
-    TorchLibOpInfo("all_dim", core_ops.aten_all_dim).xfail(
-        matcher=lambda sample: not (len(sample.kwargs) > 0),
-        reason="this Aten overload only support one tensor as input and {dim,keepdim} as kwargs by design",
+    TorchLibOpInfo("all_dim", core_ops.aten_all_dim).skip(
+        matcher=lambda sample: not (len(sample.kwargs) > 0)
+        or isinstance(sample.kwargs.get("dim"), tuple),
+        reason="this Aten overload only support one tensor as input and {dim,keepdim} as kwargs by design. dim must be an integer",
+    ),
+    TorchLibOpInfo("all_dims", core_ops.aten_all_dims, trace_only=True).skip(
+        matcher=lambda sample: not isinstance(sample.kwargs.get("dim"), tuple),
+        reason="this overload requires dim to be a tuple",
     ),
     TorchLibOpInfo("allclose", core_ops.aten_allclose),
     TorchLibOpInfo(
@@ -561,8 +566,13 @@ def _where_input_wrangler(
         "any_dim",
         core_ops.aten_any_dim,
     ).skip(
-        matcher=lambda sample: not (len(sample.kwargs) > 0),
-        reason="this Aten overload only support one tensor as input and {dim,keepdim} as kwargs by design",
+        matcher=lambda sample: not (len(sample.kwargs) > 0)
+        or isinstance(sample.kwargs.get("dim"), tuple),
+        reason="this Aten overload only support one tensor as input and {dim,keepdim} as kwargs by design. dim must be an integer",
+    ),
+    TorchLibOpInfo("any_dims", core_ops.aten_any_dims, trace_only=True).skip(
+        matcher=lambda sample: not isinstance(sample.kwargs.get("dim"), tuple),
+        reason="this overload requires dim to be a tuple",
     ),
     TorchLibOpInfo("asin", core_ops.aten_asin),
     TorchLibOpInfo("asinh", core_ops.aten_asinh),
@@ -881,7 +891,9 @@ def _where_input_wrangler(
     TorchLibOpInfo("log2", core_ops.aten_log2),
     TorchLibOpInfo("logaddexp", core_ops.aten_logaddexp),
     TorchLibOpInfo("logaddexp2", core_ops.aten_logaddexp2),
-    TorchLibOpInfo("logcumsumexp", core_ops.aten_logcumsumexp),
+    TorchLibOpInfo(
+        "logcumsumexp", core_ops.aten_logcumsumexp, tolerance={torch.float16: (1e-2, 1e-1)}
+    ),
     TorchLibOpInfo("logdet", core_ops.aten_logdet),
     TorchLibOpInfo("logsumexp", core_ops.aten_logsumexp),
     TorchLibOpInfo("lt", core_ops.aten_lt),
@@ -2080,8 +2092,8 @@ def _where_input_wrangler(
     TorchLibOpInfo("zeros_like", core_ops.aten_zeros_like, trace_only=True),
 )
 
-ops_test_common.duplicate_opinfo(OPS_DB, "all", ("all_dim",))
-ops_test_common.duplicate_opinfo(OPS_DB, "any", ("any_dim",))
+ops_test_common.duplicate_opinfo(OPS_DB, "all", ("all_dim", "all_dims"))
+ops_test_common.duplicate_opinfo(OPS_DB, "any", ("any_dim", "any_dims"))
 ops_test_common.duplicate_opinfo(OPS_DB, "arange", ("arange_start", "arange_start_step"))
 ops_test_common.duplicate_opinfo(OPS_DB, "argmax", ("argmax_dim",))
 ops_test_common.duplicate_opinfo(OPS_DB, "argmin", ("argmin_dim",))