migrate more simple gso checks

laithsakka · laithsakka · commit 512adb220cd5 · 2025-08-09T10:27:31.000-07:00
ghstack-source-id: 4a17433 Pull-Request: #160253
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
@@ -1780,9 +1780,9 @@ def _fused_rms_norm_backward(
 
     N = prod(inner_dims)  # type: ignore[arg-type]
     M = prod(outer_dims)  # type: ignore[arg-type]
-    from torch.fx.experimental.symbolic_shapes import guard_size_oblivious
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
 
-    if guard_size_oblivious(M <= 0) or guard_size_oblivious(N <= 0):
+    if guard_or_false(M == 0) or guard_or_false(N == 0):
         return (
             input.new_zeros(input_shape) if output_mask[0] else None,
             input.new_zeros(input_shape[axis:]) if output_mask[1] else None,
@@ -3987,9 +3987,9 @@ def _unsafe_masked_index(x, mask, indices, fill):
         lambda: "tensors used as masks must be bool tensors",
     )
 
-    from torch.fx.experimental.symbolic_shapes import guard_size_oblivious
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
 
-    if guard_size_oblivious(x.numel() == 0):
+    if guard_or_false(x.numel() == 0):
         meta_result = torch._meta_registrations.meta_index_Tensor(x, indices)
         return x.new_full(meta_result.shape, fill)
 
diff --git a/torch/_prims_common/__init__.py b/torch/_prims_common/__init__.py
@@ -1879,10 +1879,12 @@ def compute_required_storage_length(
     40
 
     """
-    from torch.fx.experimental.symbolic_shapes import guard_size_oblivious
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
 
     # Short-circuits if the shape has no elements
-    if guard_size_oblivious(reduce(operator.mul, shape, 1) == 0):
+    # Note: we are unsafely assuming tensor is not empty here, without
+    # runtime assertions.
+    if guard_or_false(reduce(operator.mul, shape, 1) == 0):
         return 0
 
     max_offset = sum((x - 1) * y for x, y in zip(shape, strides))
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
@@ -412,7 +412,7 @@ def _broadcast_shapes(*_shapes):
     for arg_idx, shape in enumerate(shapes):
         for idx in range(-1, -1 - len(shape), -1):
             if is_nested_int(shape[idx]):
-                # maintian nested int behaviour added in PR 145957.
+                # maintain nested int behaviour added in PR 145957.
                 if is_nested_int(common_shape[idx]) and guard_or_false(
                     shape[idx] == common_shape[idx]
                 ):
@@ -4194,7 +4194,7 @@ def index_select(x: TensorLike, dim: int, index: TensorLike):
 
 @register_decomposition(aten.squeeze.dims)
 def squeeze(a: TensorLikeType, dim: Optional[DimsType] = None) -> TensorLikeType:
-    from torch.fx.experimental.symbolic_shapes import guard_size_oblivious
+    from torch.fx.experimental.symbolic_shapes import guard_or_false
 
     if dim is None:
         dims = tuple(idx for idx, size in enumerate(a.shape) if size == 1)
@@ -4209,7 +4209,8 @@ def squeeze(a: TensorLikeType, dim: Optional[DimsType] = None) -> TensorLikeType
         return prims.view_of(a)
 
     # Note: squeeze does not modify tensors when the given dim is not a dimension of length 1
-    dims = tuple(d for d in dims if guard_size_oblivious(a.shape[d] == 1))
+    # would it be better if we just not allow 1 for unbacked at runtiume?
+    dims = tuple(d for d in dims if guard_or_false(a.shape[d] == 1))
     if len(dims) == 0:
         return prims.view_of(a)
     if len(dims) == 1:
