unify broadcast_shapes functions and avoid duplicates

laithsakka · laithsakka · commit 9e500e39678d · 2025-08-09T10:43:08.000-07:00
ghstack-source-id: 797e040 Pull-Request: #160251
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
@@ -385,7 +385,11 @@ def handle_noncontiguous_outputs(input_tlist, output):
 
 
 def _broadcast_shapes(*_shapes):
-    from torch.fx.experimental.symbolic_shapes import guard_or_false
+    from torch.fx.experimental.symbolic_shapes import (
+        guard_or_false,
+        guard_size_oblivious,
+        is_nested_int,
+    )
 
     shapes = tuple(
         (x,) if isinstance(x, IntLike) else x
@@ -407,16 +411,24 @@ def _broadcast_shapes(*_shapes):
     ] * reduce(max, (len(shape) for shape in shapes))
     for arg_idx, shape in enumerate(shapes):
         for idx in range(-1, -1 - len(shape), -1):
-            # if both 1, or statically known the same, we rather pick non-broadcast path.
-            if guard_or_false(common_shape[idx] == shape[idx]):
-                continue
-            elif guard_or_false(common_shape[idx] == 1):
+            if is_nested_int(shape[idx]):
+                # maintain nested int behaviour added in PR 145957.
+                if is_nested_int(common_shape[idx]) and guard_or_false(
+                    shape[idx] == common_shape[idx]
+                ):
+                    continue
+            else:
+                if guard_or_false(shape[idx] == common_shape[idx]):
+                    continue
+
+            if guard_or_false(common_shape[idx] == 1):
                 if shape[idx] < 0:
                     raise ValueError(
                         "Attempting to broadcast a dimension with negative length!"
                     )
                 common_shape[idx] = shape[idx]
-            elif guard_or_false(shape[idx] == 1):
+
+            if not is_nested_int(shape[idx]) and guard_or_false(shape[idx] == 1):
                 # broadcast case .
                 continue
             else:
@@ -2780,9 +2792,7 @@ def cat_compute_output_memory_format(inputs):
 
     utils.check_same_device(*tensors, allow_cpu_scalar_tensors=False)
 
-    from torch.fx.experimental.symbolic_shapes import (
-        guard_or_false,
-    )
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
 
     # This is a bit tricky.  Naively, you would expect to just pick one
     # arbitrary tensor and check that all tensors match this tensor.  However,
diff --git a/torch/functional.py b/torch/functional.py
@@ -105,58 +105,9 @@ def broadcast_shapes(*shapes):
     # This wrapper exists to support variadic args.
     # TODO Move this to C++ once the jit has better support for torch.Size.
     if not torch.jit.is_tracing():
-        max_len = 0
-        for shape in shapes:
-            if isinstance(shape, (int, torch.SymInt)):
-                if max_len < 1:
-                    max_len = 1
-            elif isinstance(shape, (tuple, list)):
-                s = len(shape)
-                if max_len < s:
-                    max_len = s
-        result = [1] * max_len
-
-        from torch.fx.experimental.symbolic_shapes import (
-            guard_size_oblivious,
-            is_nested_int,
-        )
-
-        for shape in shapes:
-            if isinstance(shape, (int, torch.SymInt)):
-                shape = (shape,)
-            if isinstance(shape, (tuple, list)):
-                for i in range(-1, -1 - len(shape), -1):
-                    if shape[i] < 0:
-                        raise RuntimeError(
-                            f"Trying to create tensor with negative dimension ({shape[i]}): ({shape[i]})"
-                        )
-
-                    # NB: handle nested ints specially to avoid invalid guarding on Ne(j0, 1).
-                    if is_nested_int(shape[i]):
-                        # Broadcasting is allowed for (j0, 1) or (j0, j0);
-                        # not (j0, j1), (j0, 5), etc.
-                        if is_nested_int(result[i]) and guard_size_oblivious(
-                            shape[i] == result[i]
-                        ):
-                            continue
-                    else:
-                        # NB: result is initialized to 1 so this is effectively an
-                        # equals one test
-                        if guard_size_oblivious(shape[i] == 1) or guard_size_oblivious(
-                            shape[i] == result[i]
-                        ):
-                            continue
-
-                    if result[i] != 1:
-                        raise RuntimeError(
-                            "Shape mismatch: objects cannot be broadcast to a single shape"
-                        )
-                    result[i] = shape[i]
-            else:
-                raise RuntimeError(
-                    "Input shapes should be of type ints, a tuple of ints, or a list of ints, got ",
-                    shape,
-                )
+        result = torch._refs._broadcast_shapes(*shapes)
+        if result is None:
+            return torch.Size([])
         return torch.Size(result)
     else:
         # with implementation above, torch.jit.trace hardcodes the sizes which makes subsequent replays fail
