pytorch
diff --git a/‎test/test_dynamic_shapes.py
+41 b/‎test/test_dynamic_shapes.py
+41
diff --git a/‎torch/fx/experimental/symbolic_shapes.py
+32-5 b/‎torch/fx/experimental/symbolic_shapes.py
+32-5
diff --git a/‎torch/fx/passes/runtime_assert.py
+4-1 b/‎torch/fx/passes/runtime_assert.py
+4-1
@@ -3005,6 +3005,15 @@ def test_remove_symbols_without_guarding(self):
         self.assertEqual(f"{x_clean.stride()}", "(8, 1)")
         self.assertEqual(f"{x_clean.shape}", "torch.Size([5, 8])")
 
+
+def custom_pass(graph: torch.fx.Graph) -> torch.fx.Graph:
+    for node in graph.nodes:
+        if node.name == "arg3_1":
+            assert node.meta["val"].size()[0] == 2
+    return graph
+
+
+class TestUnbacked(TestCase):
     @torch._dynamo.config.patch("capture_scalar_outputs", True)
     def test_deferred_neq_assert(self):
         @torch.compile(fullgraph=True)
@@ -3017,6 +3026,38 @@ def func(a):
         with self.assertRaises(RuntimeError):
             func(torch.tensor([5]))
 
+    # Test a situation where we generate a runtime assert i.e: u1==s1, then we spcialize s1
+    # later on to a constant.
+    @torch._dynamo.config.patch("capture_s
10000
calar_outputs", True)
+    def test_post_specialize_runtime_assert1(self):
+        @torch.compile(dynamic=True)
+        def func(x, y):
+            u0 = y.item()
+            s0 = x.size()[0]
+            s1 = x.size()[1]
+            torch._check(u0 + s0 + s1 == 102)
+            assert s0 == 2
+            return x * 10
+
+        func(torch.rand(2, 50), torch.tensor([50]))
+        with self.assertRaises(RuntimeError):
+            func(torch.rand(2, 50), torch.tensor([51]))
+
+    @torch._dynamo.config.patch("capture_scalar_outputs", True)
+    @torch._inductor.config.patch(post_grad_custom_pre_pass=custom_pass)
+    def test_post_specialize_runtime_assert2(self):
+        @torch.compile(dynamic=True)
+        def func(x, y):
+            u0 = y.item()
+            s0 = x.size()[0]
+            s1 = x.size()[1]
+            torch._check(u0 + s0 + s1 == 102)
+            return x * 10
+
+        func(torch.rand(2, 50), torch.tensor([50]))
+        with self.assertRaises(RuntimeError):
+            func(torch.rand(2, 50), torch.tensor([51]))
+
     @torch._dynamo.config.patch("capture_scalar_outputs", True)
     def test_deferred_sym_or_assert(self):
         @torch.compile(fullgraph=True)
 
@@ -2320,11 +2320,18 @@ def wrapper(self: ShapeEnv, *args: Any, **kwargs: Any) -> _T:  # type: ignore[mi
 # and is exclusively used for things that MUST be true (unlike guards,
 # which can evaluate False, in which case you just choose not to use
 # a particular specialization)
-@dataclass(frozen=True)
+@dataclass()
 class RuntimeAssert:
-    expr: SympyBoolean
+    _expr: SympyBoolean
     msg: str = field(repr=False)
     stack: CapturedTraceback = field(repr=False)
+    shape_env: ShapeEnv
+
+    @property
+    def expr(self) -> SympyBoolean:
+        # Whenever we access expr we want to replace specialized backed symbols with their corresponding
+        # specialized values.
+        return self.shape_env.replace(self._expr, only_specialized_backed=True)
 
 
 # Used for printing SymExprs in compile_fx
@@ -4483,6 +4490,12 @@ def is_unbacked_symint(self, symbol: sympy.Symbol) -> bool:
         """Check if a sympy symbol matches the naming convention for unbacked symbols"""
         return symbol_is_type(symbol, SymT.UNBACKED_INT)
 
+    def is_unbacked(self, symbol: sympy.Symbol) -> bool:
+        """Check if a sympy symbol matches the naming convention for unbacked symbols"""
+        return symbol_is_type(symbol, SymT.UNBACKED_INT) or symbol_is_type(
+            symbol, SymT.UNBACKED_FLOAT
+        )
+
     @record_shapeenv_event()
     def create_unbacked_symbool(self) -> SymBool:
         """Create a symbolic boolean without a hint value"""
@@ -5941,11 +5954,22 @@ def resimplify_floor_div(axioms: dict[sympy.Expr, sympy.Expr]) -> None:
         return r
 
     @_lru_cache
-    def replace(self, expr: _SympyT) -> _SympyT:
-        """Apply symbol replacements to any symbols in the given expression"""
+    def replace(self, expr: _SympyT, only_specialized_backed: bool = False) -> _SympyT:
+        """
+        Apply symbol replacements to any symbols in the given expression
+        If only_specialized_backed is set, only repalce backed symbols that
+        got specialized i.e: s0:10.
+        """
         replacements = {}
         for s in expr.free_symbols:
+            if only_specialized_backed and self.is_unbacked(s):
+                continue
+
             r = self._find(s)
+
+            if only_specialized_backed and not (isinstance(r, (int, float))):
+                continue
+
             # Micro-optimization: only do replacements if r and s are different
             # Otherwise, xreplace is not a no-op and will trigger expensive
             # assumption queries if expr has a relational node.
@@ -6380,6 +6404,9 @@ def _set_replacement(self, a: sympy.Symbol, tgt: sympy.Expr, msg: str) -> None:
                 },
             )
 
+            # when specializing to a constant, runtime assertions that uses a, need to
+            # be updated to remove a from them, since a will no longer be a graph input.
+
             for source in self.var_to_sources.get(a, []):
                 if user_tb:
                     self.specialization_stacks[source] = user_tb
@@ -7293,7 +7320,7 @@ def defer_runtime_assert(
             orig_expr = expr
             expr = canonicalize_bool_expr(expr)
             stack = CapturedTraceback.extract(skip=1)
-            ra = RuntimeAssert(expr, msg, stack)
+            ra = RuntimeAssert(expr, msg, stack, self)
             # TODO: Do this in a way that is less janky than int(s.name[1:])
             cands = sorted(
                 (s for s in expr.free_symbols if symbol_is_type(s, SymT.UNBACKED_INT)),
 
@@ -257,7 +257,10 @@ def add_runtime_asserts(ras):
                 # Convert the sympy expression into a sequence of FX
                 # nodes
                 with _set_node_metadata_hook(gm, _node_metadata_hook):
-                    res = _sympy_interp(expr_to_proxy, ra.expr).node
+                    try:
+                        res = _sympy_interp(expr_to_proxy, ra.expr).node
+                    except:
+                        raise
                     graph.call_function(
                         torch.ops.aten._assert_scalar.default,
                         # TODO: use ra.msg here, but it's pretty