test_dynamic_shapes.py

# Owner(s): ["oncall: jit"]
# ruff: noqa: F841
import contextlib
import copy
import itertools
import math
import operator
import unittest

import numpy as np
import pytest
import sympy

import torch
import torch.fx
import torch.nn.functional as F
from torch import sym_int, SymBool, SymFloat, SymInt
from torch._C import _disabled_torch_function_impl
from torch._dynamo.testing import CompileCounter, CompileCounterWithBackend
from torch._inductor.utils import fresh_cache
from torch.fx.experimental import sym_node
from torch.fx.experimental.proxy_tensor import make_fx
from torch.fx.experimental.sym_node import method_to_operator, SymNode, to_node
from torch.fx.experimental.symbolic_shapes import (
    _constrain_range_for_size,
    DimConstraints,
    DimDynamic,
    expect_true,
    guard_bool,
    guard_float,
    guard_int,
    guarding_hint_or_throw,
    GuardOnDataDependentSymNode,
    has_free_symbols,
    is_symbolic,
    ShapeEnv,
    StatelessSymbolicContext,
    statically_known_false,
    statically_known_true,
)
from torch.testing._internal.common_dtype import all_types_and
from torch.testing._internal.common_utils import (
    instantiate_parametrized_tests,
    parametrize,
    run_tests,
    skipIfTorchDynamo,
    TEST_XPU,
    TestCase,
)
from torch.testing._internal.logging_utils import logs_to_string
from torch.utils import _pytree as pytree
from torch.utils._python_dispatch import TorchDispatchMode
from torch.utils._sympy.functions import (
    CleanDiv,
    FloorDiv,
    IsNonOverlappingAndDenseIndicator,
    Mod,
)


aten = torch.ops.aten

meta_funcs = {}


def register_meta(op):
    def decorator(f):
        def add_func(op):
            meta_funcs[op] = f

        pytree.tree_map_(add_func, op)
        return f

    return decorator


@register_meta([aten.add.Tensor, aten.sub.Tensor])
def binary_meta(a, b):
    return a.new_empty(a.shape)


@register_meta(aten.cat.default)
def cat_meta(tensors, dim=0):
    concat_length = 0
    shape = tensors[0].shape
    for tensor in tensors:
        for idx, (common_length, length) in enumerate(zip(shape, tensor.shape)):
            if idx == dim:
                concat_length = concat_length + length
            else:
                if length != common_length:
                    raise AssertionError(
                        f"length mismatch: {length} != {common_length}"
                    )
    new_shape = list(shape)
    new_shape[dim] = concat_length
    return tensors[0].new_empty(new_shape)


@register_meta([aten.narrow_copy.default])
def narrow_copy_symint_meta(a, dim, start, length, **kwargs):
    shape = []
    for i, x in enumerate(a.shape):
        if i == dim:
            shape.append(length)
        else:
            shape.append(x)
    return a.new_empty(tuple(shape))


@register_meta([aten.expand.default])
def expand_symint_meta(a, size, implicit=False):
    return a.new_empty(size)


def create_contiguous(shape):
    strides = [1]
    for dim in reversed(shape[:-1]):
        strides.append(dim * strides[-1])
    return list(reversed(strides))


class FakeSymbolicTensor(torch.Tensor):
    @staticmethod
    def __new__(
        cls,
        sym_shape,
        sym_strides,
        dtype,
        layout,
        requires_grad,
        device,
        storage_offset=0,
    ):
        # TODO: this is wrong in general
        sym_stride = create_contiguous(sym_shape)
        r = torch.Tensor._make_wrapper_subclass(
            cls,
            sym_shape,
            sym_stride,
            storage_offset,
            dtype=dtype,
            layout=layout,
            requires_grad=requires_grad,
            device=device,
        )
        return r

    __torch_function__ = _disabled_torch_function_impl

    def new_empty(self, shape):
        return FakeSymbolicTensor(
            shape, None, self.dtype, self.layout, self.requires_grad, self.device
        )

    @classmethod
    def __torch_dispatch__(cls, func_overload, types, args=(), kwargs=None):
        if func_overload in meta_funcs:
            return meta_funcs[func_overload](*args, **kwargs)

        if func_overload == torch.ops.aten.new_empty.default:
            self = args[0]
            shape = args[1]
            return FakeSymbolicTensor(
                shape,
                self.stride(),
                self.dtype,
                self.layout,
                self.requires_grad,
                self.device,
            )

        raise RuntimeError(f"operator {func_overload} not supported")


def create_symbolic_tensor(name, arg, shape_env, source=None, dynamic_dims=None):
    from torch._dynamo.source import ConstantSource

    if source is None:
        source = ConstantSource(name)
    constraint_dims = [None] * arg.dim()
    if dynamic_dims is None:
        dynamic_dims = [DimDynamic.DUCK] * arg.dim()
    (
        sym_shapes,
        sym_strides,
        sym_storage_offset,
    ) = shape_env.create_symbolic_sizes_strides_storage_offset(
        arg,
        source=source,
        symbolic_context=StatelessSymbolicContext(
            dynamic_sizes=dynamic_dims, constraint_sizes=constraint_dims
        ),
    )
    return FakeSymbolicTensor(
        sym_shapes,
        sym_strides,
        arg.dtype,
        arg.layout,
        arg.requires_grad,
        arg.device,
        sym_storage_offset,
    )


def create_fake_tensor_with_dynamic_size(x, shape_env, dynamic_sizes, dynamic_strides):
    from torch._subclasses.fake_tensor import FakeTensorMode

    with FakeTensorMode(shape_env=shape_env) as fake_mode:
        return fake_mode.from_tensor(
            x,
            symbolic_context=StatelessSymbolicContext(
                dynamic_sizes=dynamic_sizes,
                dynamic_strides=dynamic_strides,
            ),
        )


def create_symtype(cls, pytype, shape_env, val, duck=True, **kwargs):
    from torch._dynamo.source import ConstantSource

    symbol = shape_env.create_symbol(
        val,
        source=ConstantSource(f"__testing_only{len(shape_env.backed_var_to_val)}"),
        dynamic_dim=DimDynamic.DUCK if duck else DimDynamic.DYNAMIC,
        constraint_dim=None,
        **kwargs,
    )
    return cls(SymNode(symbol, shape_env, pytype, hint=val))


# TODO: default duck to False
def create_symint(shape_env, i: int, duck=True, **kwargs) -> SymInt:
    return create_symtype(SymInt, int, shape_env, i, duck=duck, **kwargs)


def create_symbool(shape_env, b: bool) -> SymBool:
    return create_symtype(SymBool, bool, shape_env, b)


def create_symfloat(shape_env, f: float) -> SymFloat:
    return create_symtype(SymFloat, float, shape_env, f)


@skipIfTorchDynamo(
    "Creating ShapeEnv fails for confusing reasons (also we never expect dynamo to see code like this)"
)
class TestPySymInt(TestCase):
    def test_arith_ops(self):
        shape_env = ShapeEnv()
        symints = []
        for i in range(2, 5):
            symints.append((i, create_symint(shape_env, i)))

        ops = [
            operator.add,
            operator.sub,
            operator.floordiv,
            operator.mul,
            operator.mod,
        ]

        for op in ops:
            for args in itertools.permutations(symints, 2):
                if not isinstance(args[0][1], int) and (
                    (op != operator.mod or op != operator.floordiv) and args[1][0] != 0
                ):
                    self.assertTrue(
                        op(args[0][1], args[1][1]) == op(args[0][0], args[1][0])
                    )

    def test_reverse_arith_ops(self):
        shape_env = ShapeEnv()

        a = create_symint(shape_env, 2)
        self.assertTrue(5 // a == 5 // 2)

        a = create_symint(shape_env, 2)
        self.assertTrue(5 * a == 5 * 2)

    def test_sympify_symint(self):
        shape_env = ShapeEnv()
        a = create_symint(shape_env, 2)
        self.assertIs(sympy.sympify(a), a.node.expr)
        b = create_symfloat(shape_env, 3.0)
        self.assertIs(sympy.sympify(b), b.node.expr)
        c = create_symbool(shape_env, True)
        self.assertIs(sympy.sympify(c), c.node.expr)

    def test_roundtrip(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5, 4, 3), shape_env)

        self.assertTrue(not isinstance(x.shape[0], SymNode))
        self.assertTrue(isinstance(x.shape[0], SymInt))

        self.assertTrue(x.shape[0] == 5)
        self.assertTrue(x.shape[1] == 4)
        self.assertTrue(x.shape[2], 3)

        self.assertTrue(x.size()[0], 5)
        self.assertTrue(x.size()[1], 4)
        # Should be simplifiable to an integer.
        # Ref: https://github.com/pytorch/pytorch/pull/107492
        self.assertTrue(isinstance(x.size()[1], SymInt))
        self.assertTrue(
            isinstance(x.size()[1].node.maybe_as_int(), int)
        )  # due to guard above
        self.assertTrue(x.size()[2] == 3)

        self.assertTrue(x.size(0) == 5)
        self.assertTrue(x.size(1) == 4)
        self.assertTrue(x.size(2) == 3)
        self.assertTrue(isinstance(x.size(2), SymInt))
        self.assertTrue(isinstance(x.size(2).node.maybe_as_int(), int))

        y = create_symbolic_tensor("y", torch.randn(5, 4, 3)[1:], shape_env)
        self.assertTrue(isinstance(y.storage_offset(), SymInt))
        self.assertTrue(y.storage_offset() == 12)

    def test_binary(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5, 4, 3), shape_env)
        y = create_symbolic_tensor("y", torch.randn(5, 4, 3), shape_env)

        z = x + y
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

        # broadcasting
        y = create_symbolic_tensor("y2", torch.randn(1, 4, 1), shape_env)
        z = x + y
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

    def test_symint_args(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5, 4, 3), shape_env)
        y = create_symbolic_tensor("y", torch.randn(5, 4, 1), shape_env)
        LAST_DIM = 2
        z = x.narrow_copy(LAST_DIM, 0, y.shape[LAST_DIM])
        self.assertTrue(z.shape[2] == y.shape[2])

        # arithmetic expr with two symints
        z = x.narrow_copy(LAST_DIM, 0, x.shape[LAST_DIM] - y.shape[LAST_DIM])
        self.assertTrue(z.shape[2] == 2)

        # arithmetic expr with a symint and python int
        z = x.narrow_copy(LAST_DIM, 0, x.shape[LAST_DIM] - 1)
        self.assertTrue(z.shape[2] == 2)

    def test_symint_vargs(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5, 4, 3), shape_env)
        y = create_symbolic_tensor("y", torch.randn(1, 4, 1), shape_env)

        # varargs
        z = y.expand(x.shape[0], y.shape[1], x.shape[2])
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

        # shape list
        z = y.expand((x.shape[0], y.shape[1], x.shape[2]))
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

        # mixed python symints and ints
        z = y.expand(x.shape[0], y.shape[1], 3)
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

        # mixed python symints and ints in a list
        z = y.expand((x.shape[0], y.shape[1], 3))
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

        # mixed python symints and ints
        z = y.expand(5, y.shape[1], x.shape[2])
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

        # mixed python ints and symints in a list
        z = y.expand((5, y.shape[1], x.shape[2]))
        self.assertTrue(z.shape[0] == 5)
        self.assertTrue(z.shape[1] == 4)
        self.assertTrue(z.shape[2] == 3)

        z = y.expand((y.shape[1],))
        z = y.expand(y.shape[1])

    def test_symint_bitwise_and(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 0b1100)
        b0 = create_symint(shape_env, 0b1010)
        res_and = a0 & b0
        self.assertEqual(res_and, 0b1000)
        self.assertIsInstance(res_and, torch.SymInt, msg=type(res_and))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]), """Eq(BitwiseFn_bitwise_and(s97, s26), 8)"""
        )

        a1 = create_symint(shape_env, 3)
        b1 = create_symbool(shape_env, True)
        self.assertEqual(a1 & b1, 1)

        a2 = create_symint(shape_env, 0b1100)
        self.assertEqual(a2 & 0b1010, 0b1000)

        a3 = create_symbool(shape_env, True)
        b3 = create_symbool(shape_env, True)
        self.assertEqual(a3 & b3, True)

    def test_symint_bitwise_or(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 0b1100)
        b0 = create_symint(shape_env, 0b1010)
        res_or = a0 | b0
        self.assertEqual(res_or, 0b1110)
        self.assertIsInstance(res_or, torch.SymInt, msg=type(res_or))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]), """Eq(BitwiseFn_bitwise_or(s97, s26), 14)"""
        )

    def test_symint_bitwise_xor(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 0b1100)
        b0 = create_symint(shape_env, 0b1010)
        res_xor = a0 ^ b0
        self.assertEqual(res_xor, 0b0110)
        self.assertIsInstance(res_xor, torch.SymInt, msg=type(res_xor))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]), """Eq(BitwiseFn_bitwise_xor(s97, s26), 6)"""
        )

    def test_stride(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5, 5), shape_env)
        self.assertIsInstance(x.stride()[0], SymInt)

    def test_size_expressions(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5), shape_env)
        expand_x = x.expand(x.shape[0], x.shape[0])
        if expand_x.shape[0] > 3:
            result = expand_x + expand_x
        else:
            result = expand_x + expand_x

        gt_op, _bt, is_size_obv = shape_env.guards[-1]
        self.assertTrue(isinstance(gt_op, sympy.core.relational.StrictGreaterThan))
        self.assertTrue(str(x.shape[0]), str(gt_op.args[0]))
        self.assertTrue(str(expand_x.shape[1]), str(x.shape[0]))
        self.assertTrue(str(expand_x.shape[1]), str(result.shape[0]))
        self.assertFalse(is_size_obv)

    def test_floordiv_static(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 8)
        # This was extracted from
        # python test/inductor/test_cuda_cpp_wrapper.py -k
        # DynamicShapesCudaWrapperCudaTests.test_insignificant_strides_cuda_dynamic_shapes_cuda_wrapper
        bool(s0 % 2 == 0)
        bool(s0 % (s0 // 2) == 0)
        bool(2 * (s0 // 2) == s0)
        self.assertTrue(statically_known_true(s0 // (s0 // 2) == 2))

    def test_numel(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5), shape_env)
        self.assertIsInstance(x.numel(), torch.SymInt)
        self.assertIsInstance(torch.numel(x), torch.SymInt)

        x = torch.rand(3, 3)
        self.assertIsInstance(x.numel(), int)
        self.assertIsInstance(torch.numel(x), int)

    def test_int_to_float(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5), shape_env)
        r = torch.sym_float(x.shape[0])
        self.assertIsInstance(r, torch.SymFloat, msg=type(r))

    def test_aten_ops(self):
        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x", torch.randn(5), shape_env)
        torch.ops.aten.narrow_copy.default(x, 0, 0, x.shape[0])

        shape_env = ShapeEnv()
        x = create_symbolic_tensor("x2", torch.randn(5, 4, 3), shape_env)
        torch.ops.aten.expand.default(x, [x.shape[0], x.shape[1], x.shape[2]])

    def test_fx_trace_intlist(self):
        class CustomModule(torch.nn.Module):
            def forward(self, x):
                bs, c, h, w = x.shape
                return F.pad(x, (0, w % 2, 0, h % 2, 0, 0))

        m = CustomModule()
        x = torch.rand(1, 3, 4, 4)
        # should not TypeError: pad(): argument 'pad' (position 2) must be
        # tuple of ints, not tuple
        torch.fx.symbolic_trace(m)

    def test_meta_symint(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 2)
        r = torch.empty(a0, device="meta")
        self.assertIsInstance(r.shape[0], SymInt)

    def test_hash_size(self):
        # See issue #168254
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 2)
        r = torch.empty(a0, device="meta")
        self.assertRaises(TypeError, lambda: hash(r.shape))

    def test_guard_int(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 2)
        self.assertEqual(guard_int(a0), 2)
        self.assertExpectedInline(str(shape_env.guards[0][0]), """Eq(s97, 2)""")

    def test_sym_sum(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 2)
        s1 = create_symint(shape_env, 3)
        s2 = create_symint(shape_env, 4)
        self.assertEqual(
            (s0 + s1 + s2).node.expr, torch.sym_sum([s0, s1, s2]).node.expr
        )

    def test_prefer_deferred_runtime_assertions_over_guards(self):
        shape_env = ShapeEnv(prefer_deferred_runtime_asserts_over_guards=True)
        s0 = create_symint(shape_env, 2)
        self.assertEqual(guard_int(s0), 2)
        self.assertExpectedInline(str(shape_env.guards[0][0]), """Eq(s97, 2)""")

        shape_env = ShapeEnv(prefer_deferred_runtime_asserts_over_guards=True)
        s0 = create_symint(shape_env, 2)
        self.assertTrue(expect_true(s0 == 2))
        self.assertEqual(len(shape_env.guards), 0)
        self.assertExpectedInline(
            str([ra.expr for ra in shape_env.deferred_runtime_asserts[None]]),
            """[Eq(s97, 2)]""",
        )

    def test_sym_int(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 5)
        r = sym_int(a0)
        self.assertEqual(r, 5)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(str(shape_env.guards[0][0]), """Eq(s97, 5)""")

        a1 = create_symint(shape_env, 7)
        r = sym_int(a1 / 2)
        self.assertEqual(guard_int(r), 3)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[1][0]), """Eq(TruncToInt(IntTrueDiv(s26, 2)), 3)"""
        )

        a3 = create_symint(shape_env, 3)
        r = sym_int(2.0 * torch.sym_float(a3))
        self.assertEqual(guard_int(r), 6)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[2][0]), """Eq(TruncToInt(2.0*ToFloat(s57)), 6)"""
        )

    def test_sym_log2(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 4)
        r = torch._sym_log2(a0)
        self.assertEqual(r, 2.0)
        self.assertIsInstance(r, torch.SymFloat, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]), """Eq(OpaqueUnaryFn_log2(ToFloat(s97)), 2.0)"""
        )

    def test_sym_sqrt(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 4)
        r = torch._sym_sqrt(a0)
        self.assertEqual(r, 2)
        self.assertIsInstance(r, torch.SymFloat, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]), """Eq(OpaqueUnaryFn_sqrt(ToFloat(s97)), 2.0)"""
        )

    def test_sym_floor(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 5)
        r = math.floor(a0 / 2)
        self.assertEqual(r, 2)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]),
            """Eq(FloorToInt(IntTrueDiv(s97, 2)), 2)""",
        )
        r = math.floor(3.0 * a0)
        self.assertEqual(r, 15)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[1][0]),
            """Eq(FloorToInt(3.0*ToFloat(s97)), 15)""",
        )

    def test_sym_trunc(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 5)
        r = math.trunc(a0 / 2)
        self.assertEqual(r, 2)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]), """Eq(TruncToInt(IntTrueDiv(s97, 2)), 2)"""
        )
        r = torch.sym_int(torch.sym_sqrt(a0))
        self.assertEqual(r, 2)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[1][0]),
            """Eq(TruncToInt(OpaqueUnaryFn_sqrt(ToFloat(s97))), 2)""",
        )

    def test_sym_ceil(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 5)
        r = math.ceil(a0 / 2)
        self.assertEqual(r, 3)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]),
            """Eq(CeilToInt(IntTrueDiv(s97, 2)), 3)""",
        )
        r1 = 3.0 * a0
        r = math.floor(r1)
        self.assertEqual(r, 15)
        self.assertIsInstance(r, torch.SymInt, msg=type(r))
        self.assertExpectedInline(
            str(shape_env.guards[1][0]),
            """Eq(FloorToInt(3.0*ToFloat(s97)), 15)""",
        )

    def test_sym_ite(self):
        shape_env = ShapeEnv()
        t = create_symint(shape_env, 5)
        f = create_symint(shape_env, 4)
        b1 = True
        r1 = torch.sym_ite(b1, t, f)
        self.assertTrue(r1 is t)
        b2 = False
        r2 = torch.sym_ite(b2, t, f)
        self.assertTrue(r2 is f)
        b3 = t == 5
        r3 = torch.sym_ite(b3, t, f)
        self.assertEqual(len(shape_env.guards), 0)
        self.assertEqual(r3, 5)
        self.assertEqual(type(t), type(r3))
        self.assertExpectedInline(
            str(shape_env.guards[0][0]),
            """Eq(Piecewise((s97, Eq(s97, 5)), (s26, True)), 5)""",
        )
        b4 = f == 5
        r4 = torch.sym_ite(b4, t, f)
        self.assertEqual(len(shape_env.guards), 1)
        self.assertEqual(r4, 4)
        self.assertEqual(type(f), type(r4))
        self.assertExpectedInline(
            str(shape_env.guards[1][0]),
            """Eq(Piecewise((s97, Eq(s26, 5)), (s26, True)), 4)""",
        )

    def test_tracing_sym_ite(self):
        def f(x):
            b = x.shape[0] == 5
            ret = torch.sym_ite(b, x.shape[0], x.shape[1])
            return ret

        gm = make_fx(f, tracing_mode="symbolic")(torch.ones(4, 5))
        self.assertEqual(len(gm.shape_env.guards), 0)
        self.assertExpectedInline(
            gm.code.strip(),
            """\
def forward(self, x_1):
    sym_size_int = torch.ops.aten.sym_size.int(x_1, 0)
    eq = sym_size_int == 5
    sym_size_int_1 = torch.ops.aten.sym_size.int(x_1, 1);  x_1 = None
    sym_ite = torch.sym_ite(eq, sym_size_int, sym_size_int_1);  eq = sym_size_int = sym_size_int_1 = None
    return sym_ite""",
        )
        r1 = gm(torch.ones(4, 5))
        self.assertIsInstance(r1, int)
        self.assertEqual(r1, 5)
        r2 = gm(torch.ones(5, 4))
        self.assertIsInstance(r2, int)
        self.assertEqual(r2, 5)

    def test_int_conversion(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 2)
        int(a0)
        self.assertExpectedInline(str(shape_env.guards[0][0]), """Eq(s97, 2)""")

    def test_data_dependent_guard(self):
        shape_env = ShapeEnv()
        s0 = shape_env.create_unbacked_symint()
        self.assertRaises(GuardOnDataDependentSymNode, lambda: bool(s0 == 0))

    def test_data_dependent_guard_propagate_real_tensors(self):
        shape_env = ShapeEnv()
        s0 = shape_env.create_unbacked_symint()
        shape_env.set_real_tensor_prop_unbacked_vals(s0.node.expr, 0)
        self.assertEqual(bool(s0 == 0), True)

    def test_expect_true_basic(self):
        shape_env = ShapeEnv()
        i0 = shape_env.create_unbacked_symint()
        i0_sym = i0.node.expr
        # This doesn't error
        self.assertTrue(expect_true(i0 == 0))
        # This generates a deferred runtime assert via replacement
        self.assertEqual(shape_env.replacements[i0_sym], 0)
        # After expecting true, guards now resolve given the runtime assert
        bool(i0 == 0)

    def test_expect_true_with_s0(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 5)
        i0 = shape_env.create_unbacked_symint()
        self.assertTrue(expect_true(i0 < s0))
        self.assertExpectedInline(
            str([ra.expr for ra in shape_env.deferred_runtime_asserts[i0.node.expr]]),
            """[u0 < s97]""",
        )
        self.assertTrue(i0 < s0)
        self.assertTrue(i0 != s0)
        self.assertFalse(i0 > s0)
        self.assertFalse(i0 >= s0)

    def test_expect_true_prefer_later(self):
        shape_env = ShapeEnv()
        i0 = shape_env.create_unbacked_symint()
        i1 = shape_env.create_unbacked_symint()
        i1_sym = i1.node.expr
        self.assertTrue(expect_true(i0 + i1 == 10))
        # Importantly, this is put in i1, not i0!
        self.assertExpectedInline(
            str([ra.expr for ra in shape_env.deferred_runtime_asserts[i1_sym]]),
            """[Eq(u0 + u1, 10)]""",
        )
        self.assertTrue(i0 + i1 == 10)
        # NB: We currently don't support deriving that we can substitute
        # i0 + i1 with 10; maybe we should, but this means our rewriting
        # system is no longer confluent (it's probably OK though, because
        # you're unlikely to get other equalities like this on the
        # unbacked SymInts.)

    def test_unbacked_substitution(self):
        shape_env = ShapeEnv()
        i0 = shape_env.create_unbacked_symint()
        i1 = shape_env.create_unbacked_symint()
        _constrain_range_for_size(i0)
        _constrain_range_for_size(i1)
        self.assertTrue(expect_true(i0 == i1 * 4))
        self.assertExpectedInline(str(i0), """u0""")

        i2 = shape_env.create_unbacked_symint()
        i3 = shape_env.create_unbacked_symint()
        _constrain_range_for_size(i2)
        _constrain_range_for_size(i3)
        self.assertTrue(expect_true(i2 * 4 == i3))
        self.assertExpectedInline(str(i3), """u3""")

    def test_avoid_unbacked_substitution(self):
        shape_env = ShapeEnv()
        i0 = shape_env.create_unbacked_symint()
        _constrain_range_for_size(i0)
        i1 = shape_env.create_unbacked_symint()
        _constrain_range_for_size(i1)
        self.assertTrue(expect_true(i0 == 10 - i1))
        self.assertExpectedInline(str(i0), """u0""")

    def test_expect_true_double_digits(self):
        shape_env = ShapeEnv()
        ia = [shape_env.create_unbacked_symint() for _ in range(11)]  # allocate 10
        self.assertEqual(str(ia[-1]), "u10")
        self.assertTrue(expect_true(sum(ia) == 20))
        self.assertEqual(len(shape_env.deferred_runtime_asserts[ia[-1].node.expr]), 1)

    def test_expect_true_refine_range(self):
        shape_env = ShapeEnv()
        for i, rel in enumerate(
            [lambda x: x > 4, lambda x: 4 < x, lambda x: x >= 5, lambda x: 5 <= x]
        ):
            with self.subTest(f"i = {i}"):
                i0 = shape_env.create_unbacked_symint()
                self.assertTrue(expect_true(rel(i0)))
                self.assertTrue(statically_known_true(i0 != 3))
                self.assertTrue(statically_known_true(i0 != 4))
                self.assertFalse(statically_known_true(i0 != 5))
                self.assertFalse(statically_known_true(i0 != 6))
                self.assertTrue(statically_known_true(i0 > 4))
                self.assertTrue(statically_known_true(i0 >= 5))

        for i, rel in enumerate(
            [lambda x: x < 4, lambda x: 4 > x, lambda x: x <= 3, lambda x: 3 >= x]
        ):
            with self.subTest(f"i = {i}"):
                i0 = shape_env.create_unbacked_symint()
                self.assertTrue(expect_true(rel(i0)))
                self.assertFalse(statically_known_true(i0 != 2))
                self.assertFalse(statically_known_true(i0 != 3))
                self.assertTrue(statically_known_true(i0 != 4))
                self.assertTrue(statically_known_true(i0 != 5))
                self.assertTrue(statically_known_true(i0 < 4))
                self.assertTrue(statically_known_true(i0 <= 5))

    def test_guard_refine_range(self):
        shape_env = ShapeEnv()
        for i, rel in enumerate(
            [lambda x: x > 4, lambda x: 4 < x, lambda x: x >= 5, lambda x: 5 <= x]
        ):
            with self.subTest(f"i = {i}"):
                i0 = create_symint(shape_env, 10, duck=False)
                self.assertTrue(bool(rel(i0)))
                self.assertTrue(statically_known_true(i0 != 3))
                self.assertTrue(statically_known_true(i0 != 4))
                self.assertFalse(statically_known_true(i0 != 5))
                self.assertFalse(statically_known_true(i0 != 6))
                self.assertTrue(statically_known_true(i0 > 4))
                self.assertTrue(statically_known_true(i0 >= 5))

        for i, rel in enumerate(
            [lambda x: x > 4, lambda x: 4 < x, lambda x: x >= 5, lambda x: 5 <= x]
        ):
            with self.subTest(f"i = {i}"):
                i0 = create_symint(shape_env, 2, duck=False)
                self.assertFalse(bool(rel(i0)))
                self.assertFalse(statically_known_true(i0 != 3))
                self.assertFalse(statically_known_true(i0 != 4))
                self.assertTrue(statically_known_true(i0 != 5))
                self.assertTrue(statically_known_true(i0 != 6))
                self.assertTrue(statically_known_true(i0 <= 4))
                self.assertTrue(statically_known_true(i0 < 5))

        for i, rel in enumerate(
            [lambda x: x < 4, lambda x: 4 > x, lambda x: x <= 3, lambda x: 3 >= x]
        ):
            with self.subTest(f"i = {i}"):
                i0 = create_symint(shape_env, 2, duck=False)
                self.assertTrue(bool(rel(i0)))
                self.assertFalse(statically_known_true(i0 != 2))
                self.assertFalse(statically_known_true(i0 != 3))
                self.assertTrue(statically_known_true(i0 != 4))
                self.assertTrue(statically_known_true(i0 != 5))
                self.assertTrue(statically_known_true(i0 < 4))
                self.assertTrue(statically_known_true(i0 <= 3))

        for i, rel in enumerate(
            [lambda x: x < 4, lambda x: 4 > x, lambda x: x <= 3, lambda x: 3 >= x]
        ):
            with self.subTest(f"i = {i}"):
                i0 = create_symint(shape_env, 10, duck=False)
                self.assertFalse(bool(rel(i0)))
                self.assertTrue(statically_known_true(i0 != 2))
                self.assertTrue(statically_known_true(i0 != 3))
                self.assertFalse(statically_known_true(i0 != 4))
                self.assertFalse(statically_known_true(i0 != 5))
                self.assertTrue(statically_known_true(i0 >= 4))
                self.assertTrue(statically_known_true(i0 > 3))

    def test_mul_int_oo_nan(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 5, duck=False)
        s1 = create_symint(shape_env, 6, duck=False)
        s2 = create_symint(shape_env, 5, duck=False)
        bool(s0 * (s1 // s0) == s2)

    def test_non_overlapping_and_dense_backed(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 5)
        r = torch.empty_strided((a0, 7), (1, a0), device="meta")
        self.assertTrue(torch.ops.aten.is_non_overlapping_and_dense.default(r))

    def test_non_overlapping_and_dense_unbacked(self):
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        cf = torch.ops.aten.is_non_overlapping_and_dense.default

        self.assertEqual(IsNonOverlappingAndDenseIndicator(u0.node.expr, 2, 2, 1), 1)
        self.assertEqual(IsNonOverlappingAndDenseIndicator(2, u0.node.expr, 1, 2), 1)
        self.assertTrue(cf(torch.empty_strided((u0, 2), (2, 1), device="meta")))
        self.assertTrue(cf(torch.empty_strided((2, u0), (1, 2), device="meta")))

        self.assertEqual(IsNonOverlappingAndDenseIndicator(u0.node.expr, 1), 1)
        self.assertEqual(IsNonOverlappingAndDenseIndicator(1, u0.node.expr), 1)
        self.assertTrue(cf(torch.empty_strided((u0,), (1,), device="meta")))
        self.assertTrue(cf(torch.empty_strided((1,), (u0,), device="meta")))

        Max = torch.sym_max
        # NB: This only works because we're able to determine this tensor is
        # contiguous. transpose(0, 1) makes it stop working
        self.assertTrue(
            cf(
                torch.empty_strided(
                    (2, 3, 1, u0),
                    (3 * Max(1, u0), Max(1, u0), Max(1, u0), 1),
                    device="meta",
                )
            )
        )

    def test_prims_is_non_overlapping_and_dense_or_false(self):
        shape_env = ShapeEnv()
        cf = torch._prims_common.is_non_overlapping_and_dense_or_false

        # backed case
        a0 = create_symint(shape_env, 5)
        self.assertTrue(cf(torch.empty_strided((a0, 7), (1, a0), device="meta")))

        # unbacked
        u0 = shape_env.create_unbacked_symint()
        self.assertTrue(cf(torch.empty_strided((u0, 2), (2, 1), device="meta")))
        self.assertTrue(cf(torch.empty_strided((2, u0), (1, 2), device="meta")))
        self.assertTrue(cf(torch.empty_strided((u0,), (1,), device="meta")))
        self.assertTrue(cf(torch.empty_strided((1,), (u0,), device="meta")))

        Max = torch.sym_max
        self.assertTrue(
            cf(
                torch.empty_strided(
                    (2, 3, 1, u0),
                    (3 * Max(1, u0), Max(1, u0), Max(1, u0), 1),
                    device="meta",
                )
            )
        )
        self.assertFalse(
            cf(
                torch.empty_strided(
                    (2, 3, 1, u0),
                    (Max(1, u0), Max(1, u0), 1, 3 * Max(1, u0)),
                    device="meta",
                )
            )
        )

        # return False on arbitrary strides
        u1 = shape_env.create_unbacked_symint()
        self.assertFalse(
            cf(
                torch.empty_strided(
                    (2 * u0, u0, 1),
                    (u1, u0, u0 + u1),
                    device="meta",
                )
            )
        )
        self.assertFalse(
            cf(
                torch.empty_strided(
                    (2, 3, u0),
                    (u1, 3, 1),
                    device="meta",
                )
            )
        )

    def test_sympy_optimized_add_binary_search(self):
        import sympy

        from torch.fx.experimental.sym_node import _binary_search_insert_arg

        a = sympy.Symbol("a")
        b = sympy.Symbol("b")
        c = sympy.Symbol("c")

        args = []
        args = _binary_search_insert_arg([], b)
        self.assertEqual(args, [b])

        self.assertEqual(_binary_search_insert_arg(args, b), None)

        args = _binary_search_insert_arg(args, a)
        self.assertEqual(args, [a, b])

        self.assertEqual(_binary_search_insert_arg(args, b), None)
        self.assertEqual(_binary_search_insert_arg(args, a), None)

        args = _binary_search_insert_arg(args, c)
        self.assertEqual(args, [a, b, c])

        self.assertEqual(_binary_search_insert_arg(args, a), None)
        self.assertEqual(_binary_search_insert_arg(args, b), None)
        self.assertEqual(_binary_search_insert_arg(args, c), None)

        a1 = sympy.Symbol("a1")
        a2 = sympy.Symbol("a2")

        args = _binary_search_insert_arg(args, a1)
        self.assertEqual(args, [a, a1, b, c])

        args = _binary_search_insert_arg(args, a2)
        self.assertEqual(args, [a, a1, a2, b, c])

        c1 = sympy.Symbol("c1")
        args = _binary_search_insert_arg(args, c1)
        self.assertEqual(args, [a, a1, a2, b, c, c1])

        # insert to front
        _a = sympy.Symbol("_a")
        args = _binary_search_insert_arg(args, _a)
        self.assertEqual(args, [_a, a, a1, a2, b, c, c1])

    def test_floor_clean_div_axioms(self):
        # Test that if we add an axiom that have FloorDiv, after which the
        # shapeEnv changed such that it can be simplified it to CleanDiv, then
        # We still correctly replace CleanDiv with the axiom value of FloorDiv.
        shape_env = ShapeEnv()
        a = shape_env.create_unbacked_symint()

        shape_env.guard_or_defer_runtime_assert((a // 3 == 1).node.expr, " test")

        from sympy import Eq

        test1 = Eq(FloorDiv(a.node.expr, 3), 1)
        test2 = Eq(CleanDiv(a.node.expr, 3), 1)

        self.assertTrue(shape_env.evaluate_expr(test1))
        self.assertEqual(shape_env._maybe_evaluate_static(test2), None)

        # After this FloorDiv(a, 3) is simplified to CleanDiv(a, 3)
        shape_env.guard_or_defer_runtime_assert(Eq(Mod(a, 3), 0), " test")
        self.assertEqual(test2, shape_env.simplify(test1))

        self.assertTrue(shape_env.evaluate_expr(test1))
        self.assertTrue(shape_env.evaluate_expr(test2))

    def test_sympy_optimized_add(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 2)
        s1 = create_symint(shape_env, 3)
        s2 = create_symint(shape_env, 4)
        sum = s0 + s1

        self.assertTrue(sum.node._optimized_summation)

        def assert_optimized(sym):
            self.assertTrue(sym.node._optimized_summation)

        def assert_not_optimized(sym):
            self.assertFalse(getattr(sym.node, "_optimized_summation", False))

        assert_optimized(sum)

        # add duplicate symbol
        assert_not_optimized(sum + s0)

        # add constant.
        assert_not_optimized(sum + 1)

        # add new unique symbol, should maintain _optimized_summation property.
        assert_optimized(sum + s2)

        assert_optimized(s2 + sum)

        # add x + (a+b) with no  _optimized_summation on the rhs sum.
        a = create_symint(shape_env, 10)
        b = create_symint(shape_env, 11)
        two_sum = torch.sym_sum([a, b])
        assert_not_optimized(two_sum)
        assert_optimized(sum + two_sum)

        # adding two expressions of length >2 that are _optimized_summation.
        a = s0 + s1 + s2
        s3 = create_symint(shape_env, 10)
        s4 = create_symint(shape_env, 20)
        s5 = create_symint(shape_env, 30)
        b = s3 + s4 + s5
        assert_optimized(a)
        assert_optimized(b)
        assert_not_optimized(a + b)
        assert_not_optimized(b + a)
        assert_not_optimized(b + a + b)

    def test_max_of_unique_summation_opt(self):
        shape_env = ShapeEnv()
        s0 = shape_env.create_unbacked_symint()
        s1 = shape_env.create_unbacked_symint()
        s2 = shape_env.create_unbacked_symint()
        s3 = shape_env.create_unbacked_symint()
        s4 = shape_env.create_unbacked_symint()
        s5 = shape_env.create_unbacked_symint()
        s7 = shape_env.create_unbacked_symint()

        def assert_optimized(sym):
            self.assertTrue(sym.node.expr.unique_summations_symbols is not None)

        def assert_not_optimized(sym):
            getattr(sym.node.expr, "unique_summations_symbols", None)

        mx1 = torch.sym_max(s0, s1)
        assert_not_optimized(mx1)

        mx2 = torch.sym_max(s0 + s1, s2 + s3)
        assert_optimized(mx2)

        mx3 = torch.sym_max(mx2, s4 + s5)
        assert_optimized(mx3)
        assert_optimized(torch.sym_max(s4 + s5, mx2))

        assert_not_optimized(torch.sym_max(mx3, s7))
        assert_not_optimized(torch.sym_max(mx3, 10))
        assert_not_optimized(torch.sym_max(mx3, s3 + s7))
        assert_not_optimized(torch.sym_max(mx3, s7 * 2))

    def test_build_proxy_for_optimized_add(self):
        """
        Test that _build_proxy_for_sym_expr correctly handles flattened
        sympy.Add expressions with more than 2 arguments.

        When sympy.Add operations are optimized by _optimized_add
        (in torch/fx/experimental/sym_node.py), they can be flattened into a
        single Add node with 3 or more arguments (e.g., Add(s0, s1, s2)
        instead of Add(Add(s0, s1), s2)).

        _build_proxy_for_sym_expr must handle this by using torch.sym_sum
        (which natively supports n-ary addition) instead of operator.add
        (which is binary and would either raise TypeError or silently drop
        arguments via FX codegen's 2-placeholder format string).
        """
        import torch.fx as fx
        from torch.fx.experimental.proxy_tensor import (
            _build_proxy_for_sym_expr,
            _SympyExprTrackerValue,
            PythonKeyTracer,
            set_meta,
        )
        from torch.utils._thunk import Thunk

        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 2)
        s1 = create_symint(shape_env, 3)
        s2 = create_symint(shape_env, 4)

        # This triggers _optimized_add and produces a flattened sympy.Add
        # with 3 arguments: Add(s0_symbol, s1_symbol, s2_symbol)
        optimized_sum = s0 + s1 + s2
        self.assertTrue(optimized_sum.node._optimized_summation)
        # Confirm the underlying sympy expression has 3 args (flattened).
        self.assertEqual(len(optimized_sum.node.expr.args), 3)

        # ------------------------------------------------------------------
        # Case 1: out=None
        #
        # Previously this would call operator.add(a, b, c) and raise
        # TypeError. Now uses torch.sym_sum which handles n-ary addition.
        # ------------------------------------------------------------------
        tracer_none = PythonKeyTracer()
        for sym in [s0, s1, s2]:
            tracer_none.sympy_expr_tracker[sym.node.expr] = _SympyExprTrackerValue(
                proxy=sym, value=sym
            )

        result = _build_proxy_for_sym_expr(tracer_none, optimized_sum.node.expr)
        self.assertEqual(int(result), 9)

        # ------------------------------------------------------------------
        # Case 2: out=optimized_sum  (the typical path from get_proxy_slot)
        #
        # Previously _sym_register would record operator.add with 3 args,
        # and FX codegen's "{} + {}" format string would silently drop the
        # third argument. Now uses torch.sym_sum which correctly handles
        # all arguments.
        # ------------------------------------------------------------------
        tracer = PythonKeyTracer()
        tracer.root = torch.nn.Module()
        tracer.graph = fx.Graph(tracer_cls=PythonKeyTracer)

        for sym, name in [(s0, "s0"), (s1, "s1"), (s2, "s2")]:
            node = tracer.graph.placeholder(name)
            proxy = fx.Proxy(node, tracer)
            set_meta(proxy, sym)
            tracer.sympy_expr_tracker[sym.node.expr] = _SympyExprTrackerValue(
                proxy=proxy, value=sym
            )
            tracer.symnode_tracker[sym] = Thunk(lambda p=proxy: p)

        _build_proxy_for_sym_expr(tracer, optimized_sum.node.expr, out=optimized_sum)

        out_proxy = tracer.symnode_tracker[optimized_sum].force()
        tracer.graph.output(out_proxy.node)

        gm = fx.GraphModule(tracer.root, tracer.graph)

        result = gm(2, 3, 4)  # should be 2 + 3 + 4 = 9
        self.assertEqual(result, 9)

    def test_sym_max_multi_max_simplify(self):
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        self.assertTrue(
            statically_known_true(
                torch.sym_max(1, torch.sym_max(257, u0)) == torch.sym_max(257, u0)
            )
        )

    def test_numpy_sym_max(self):
        self.assertEqual(torch.sym_max(np.int64(10), 12), 12)
        self.assertEqual(torch.sym_max(np.int64(12), 10), 12)
        self.assertEqual(torch.sym_max(np.int64(10), 12.5), 12.5)
        self.assertEqual(torch.sym_max(np.int64(14), 12.5), 14.0)
        self.assertEqual(torch.sym_max(np.float64(14.0), 12), 14.0)
        self.assertEqual(torch.sym_max(np.float64(14.0), 16), 16.0)

    def test_numpy_sym_min(self):
        self.assertEqual(torch.sym_min(np.int64(10), 12), 10)
        self.assertEqual(torch.sym_min(np.int64(12), 10), 10)
        self.assertEqual(torch.sym_min(np.int64(10), 12.5), 10.0)
        self.assertEqual(torch.sym_min(np.int64(14), 12.5), 12.5)
        self.assertEqual(torch.sym_min(np.float64(14.0), 12), 12.0)
        self.assertEqual(torch.sym_min(np.float64(14.0), 16), 14.0)

    def test_debug_has_internal_overlap_unbacked(self):
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        cf = torch._debug_has_internal_overlap
        self.assertEqual(cf(torch.empty_strided((u0, 2), (2, 1), device="meta")), 0)
        self.assertEqual(cf(torch.empty_strided((2, u0), (1, 2), device="meta")), 0)
        self.assertEqual(cf(torch.empty_strided((u0,), (1,), device="meta")), 0)
        self.assertEqual(cf(torch.empty_strided((1,), (u0,), device="meta")), 2)
        Max = torch.sym_max
        self.assertEqual(
            cf(
                torch.empty_strided(
                    (2, 3, 1, u0),
                    (3 * Max(1, u0), Max(1, u0), Max(1, u0), 1),
                    device="meta",
                )
            ),
            2,
        )

        # Wobbling these to zero is OK too
        self.assertEqual(cf(torch.empty_strided((u0, 2), (3, 1), device="meta")), 2)
        self.assertEqual(cf(torch.empty_strided((2, u0), (1, 3), device="meta")), 2)

    def test_specialize_zero_one(self):
        shape_env = ShapeEnv(specialize_zero_one=True)
        a0 = create_symint(shape_env, 5)
        if a0 == 1:
            raise AssertionError("expected a0 != 1")
        self.assertEqual(len(shape_env.guards), 0)

        shape_env = ShapeEnv(specialize_zero_one=False)
        a0 = create_symint(shape_env, 5)
        if a0 == 1:
            raise AssertionError("expected a0 != 1")
        self.assertEqual(len(shape_env.guards), 1)

    def test_duck_shape(self):
        shape_env = ShapeEnv(duck_shape=True)
        a0 = create_symint(shape_env, 5)
        a1 = create_symint(shape_env, 5)
        if a0 != a1:
            raise AssertionError("expected a0 == a1")
        self.assertEqual(len(shape_env.guards), 0)

        shape_env = ShapeEnv(duck_shape=False)
        a0 = create_symint(shape_env, 5)
        a1 = create_symint(shape_env, 5)
        if a0 != a1:
            raise AssertionError("expected a0 == a1")
        self.assertEqual(len(shape_env.guards), 1)

    def test_int_bool(self):
        # See https://github.com/pytorch/pytorch/issues/95981
        shape_env = ShapeEnv(duck_shape=True)
        a0 = create_symint(shape_env, 5)
        if not a0:
            raise AssertionError("expected a0 to be truthy")
        self.assertEqual(len(shape_env.guards), 0)

    def test_symint_as_scalar(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 2)

        sym_int_encountered = False

        class TestSymInt(TorchDispatchMode):
            def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                if func != torch.ops.aten.add.Tensor:
                    raise AssertionError(
                        f"expected torch.ops.aten.add.Tensor, got {func}"
                    )

                nonlocal sym_int_encountered
                # WARNING: do not do identity tests on the outer
                # SymInt/SymFloat, they are NOT STABLE
                sym_int_encountered = kwargs["alpha"].node is a0.node
                kwargs["alpha"] = 0
                return func(*args)

        x = torch.rand([4, 4])
        with TestSymInt():
            y = torch.add(x, x, alpha=a0)

        self.assertTrue(sym_int_encountered)

    def test_deepcopy(self):
        shape_env = ShapeEnv()
        a0 = create_symint(shape_env, 2)
        if not (a0 < 4):
            raise AssertionError(f"expected a0 < 4, got a0={a0}")
        new_shape_env = copy.deepcopy(shape_env)
        self.assertEqual(len(new_shape_env.guards), 1)

    def test_print_readable_with_symints(self):
        def f(a, b):
            dim0 = a.shape[0] + b.shape[0]
            dim1 = a.shape[1] + b.shape[1]
            d = a.new_empty(dim0, dim1)
            d = torch.ops.aten.native_dropout(d, 0.5, train=True)
            return d

        fx_g = make_fx(f, tracing_mode="symbolic")(torch.randn(5, 3), torch.randn(4, 3))
        out = fx_g.print_readable(print_output=False)

        self.assertExpectedInline(
            out.strip(),
            """\
class f(torch.nn.Module):
    def forward(self, a_1: "f32[s75, s96]", b_1: "f32[s57, s96]"):
        # No stacktrace found for following nodes
        sym_size_int: "Sym(s75)" = torch.ops.aten.sym_size.int(a_1, 0)
        sym_size_int_1: "Sym(s57)" = torch.ops.aten.sym_size.int(b_1, 0)
        add: "Sym(s57 + s75)" = sym_size_int + sym_size_int_1;  sym_size_int = sym_size_int_1 = None
        sym_size_int_2: "Sym(s96)" = torch.ops.aten.sym_size.int(a_1, 1)
        sym_size_int_3: "Sym(s96)" = torch.ops.aten.sym_size.int(b_1, 1);  b_1 = None
        add_1: "Sym(2*s96)" = sym_size_int_2 + sym_size_int_3;  sym_size_int_2 = sym_size_int_3 = None
        new_empty: "f32[s57 + s75, 2*s96]" = torch.ops.aten.new_empty.default(a_1, [add, add_1], pin_memory = False);  a_1 = add = add_1 = None
        native_dropout = torch.ops.aten.native_dropout.default(new_empty, 0.5, True);  new_empty = None
        getitem: "f32[s57 + s75, 2*s96]" = native_dropout[0]
        getitem_1: "b8[s57 + s75, 2*s96]" = native_dropout[1];  native_dropout = None
        return (getitem, getitem_1)""",  # noqa: B950
        )

    def test_statically_known_true(self):
        shape_env = ShapeEnv()
        s2, s3, s4 = (create_symint(shape_env, i) for i in range(2, 5))

        # Statically known true
        self.assertTrue(statically_known_true(True))
        self.assertTrue(statically_known_true(s2 == s2))
        self.assertTrue(statically_known_true(s2 * s3 > s3))
        self.assertTrue(statically_known_true(s3 * s4 > s4))
        self.assertTrue(statically_known_true((s3 + s3) % 2 == 0))

        # Statically known false
        self.assertFalse(statically_known_true(False))
        self.assertFalse(statically_known_true(s3 * s4 <= s4))
        self.assertFalse(statically_known_true((s3 + s3) % 2 == 1))

        # True for hints, but not known statically
        self.assertFalse(statically_known_true(s2 + s2 == s4))
        self.assertFalse(statically_known_true(s4 % s2 == 0))
        self.assertFalse(statically_known_true(s2 != s3))
        self.assertFalse(statically_known_true(s3 * s4 > s2))

        # False for hints, but not known statically
        self.assertFalse(statically_known_true(s2 == s3))
        self.assertFalse(statically_known_true(s2 > s3))
        self.assertFalse(statically_known_true(s3 + s3 == s4))

        # No guards should be generated
        self.assertEqual(len(shape_env.guards), 0)

    def test_statically_known_false(self):
        shape_env = ShapeEnv()
        s2, s3, s4 = (create_symint(shape_env, i) for i in range(2, 5))

        # Statically known true
        self.assertFalse(statically_known_false(True))
        self.assertFalse(statically_known_false(s2 == s2))
        self.assertFalse(statically_known_false(s2 * s3 > s3))
        self.assertFalse(statically_known_false(s3 * s4 > s4))
        self.assertFalse(statically_known_false((s3 + s3) % 2 == 0))

        # Statically known false
        self.assertTrue(statically_known_false(False))
        self.assertTrue(statically_known_false(s3 * s4 <= s4))
        self.assertTrue(statically_known_false((s3 + s3) % 2 == 1))

        # True for hints, but not known statically
        self.assertFalse(statically_known_false(s2 + s2 == s4))
        self.assertFalse(statically_known_false(s4 % s2 == 0))
        self.assertFalse(statically_known_false(s2 != s3))
        self.assertFalse(statically_known_false(s3 * s4 > s2))

        # False for hints, but not known statically
        self.assertFalse(statically_known_false(s2 == s3))
        self.assertFalse(statically_known_false(s2 > s3))
        self.assertFalse(statically_known_false(s3 + s3 == s4))

        # No guards should be generated
        self.assertEqual(len(shape_env.guards), 0)

    def test_ephemeral_source_simplification(self):
        from torch._dynamo.source import EphemeralSource

        # For full robustness, ensure the ephemeral source symbols are simplified out regardless
        # of construction order or check order.
        for construct_ephemeral_first, x_first_in_check in itertools.product(
            [False, True], [False, True]
        ):
            shape_env = ShapeEnv()
            shape = (5, 10)
            dynamic_dims = [DimDynamic.DYNAMIC for _ in shape]
            x = create_symbolic_tensor(
                "x",
                torch.randn(*shape),
                shape_env,
                source=(EphemeralSource() if construct_ephemeral_first else None),
                dynamic_dims=dynamic_dims,
            )
            y = create_symbolic_tensor(
                "y",
                torch.randn(*shape),
                shape_env,
                source=(EphemeralSource() if not construct_ephemeral_first else None),
                dynamic_dims=dynamic_dims,
            )
            t_with_ephemeral = x if construct_ephemeral_first else y

            def _get_ephemeral_source_symbols(t):
                return [
                    s.node.expr
                    for s in itertools.chain(t.shape, t.stride(), (t.storage_offset(),))
                    if isinstance(s, torch.SymInt)
                    and s.node.expr in shape_env.var_to_sources
                    and any(
                        source.is_ephemeral()
                        for source in shape_env.var_to_sources[s.node.expr]
                    )
                ]

            # these checks should simplify out the ephemeral symbols, regardless of the
            # ordering x == y or y == x
            self.assertTrue(len(_get_ephemeral_source_symbols(t_with_ephemeral)) > 0)
            if x_first_in_check:
                torch._check(x.size() == y.size())
                torch._check(x.stride() == y.stride())
                torch._check(x.storage_offset() == y.storage_offset())
            else:
                torch._check(y.size() == x.size())
                torch._check(y.stride() == x.stride())
                torch._check(y.storage_offset() == x.storage_offset())
            self.assertEqual(len(_get_ephemeral_source_symbols(t_with_ephemeral)), 0)

    def test_ephemeral_source_unified_with_non_ephemeral_source(self):
        from torch._dynamo.source import EphemeralSource

        for construct_ephemeral_first in (False, True):
            shape_env = ShapeEnv()
            shape = (5, 10)
            # use duck sizing here to ensure symbol reuse across x and y
            duck_dims = [DimDynamic.DUCK for _ in shape]
            x = create_symbolic_tensor(
                "x",
                torch.randn(*shape),
                shape_env,
                source=(EphemeralSource() if construct_ephemeral_first else None),
                dynamic_dims=duck_dims,
            )
            y = create_symbolic_tensor(
                "y",
                torch.randn(*shape),
                shape_env,
                source=(EphemeralSource() if not construct_ephemeral_first else None),
                dynamic_dims=duck_dims,
            )

            # regardless of construction order, non-ephemeral sources should be preferred
            # first in the var_to_sources list for potential guarding later on
            for source_list in shape_env.var_to_sources.values():
                self.assertFalse(source_list[0].is_ephemeral())

            self.assertEqual(x.size(), y.size())
            self.assertEqual(x.stride(), y.stride())
            self.assertEqual(x.storage_offset(), y.storage_offset())

    def test_tensor_factory_with_symint(self):
        args = list(range(3))
        expected = torch.tensor(args)

        shape_env = ShapeEnv()
        sym_args = [create_symint(shape_env, i) for i in args]

        # test tensor factories
        for dt in all_types_and(torch.half, torch.bfloat16):
            res = torch.tensor(sym_args, dtype=dt)
            self.assertEqual(res, expected, exact_dtype=False)

        # test legacy tensor factories
        legacy_ctors = [
            torch.Tensor,
            torch.LongTensor,
            torch.DoubleTensor,
            torch.FloatTensor,
            torch.IntTensor,
            torch.ShortTensor,
            torch.HalfTensor,
            torch.ByteTensor,
        ]
        for Tensor in legacy_ctors:
            res = Tensor(sym_args)
            self.assertEqual(res, expected, exact_dtype=False)

    def test_backed_size_oblivious_01_spec(self):
        from torch.fx.experimental.symbolic_shapes import guard_size_oblivious

        @torch.compile(dynamic=True, fullgraph=True)
        def f(a, b):
            if guard_size_oblivious(a.size(0) == 1):
                return b * 10
            else:
                return b * 20

        with torch.fx.experimental._config.patch(backed_size_oblivious=True):
            # always go to the >= 2 branch.
            self.assertEqual(
                f(torch.tensor([1]), torch.tensor([1])), torch.tensor([20])
            )

    @fresh_cache()
    def test_slice_backed_size_oblivious(self):
        @torch.compile(backend="inductor", fullgraph=True, dynamic=True)
        def f(x):
            return x[:5]

        with torch.fx.experimental._config.patch(backed_size_oblivious=True):
            f(torch.randn(10, 10))

    def test_baddbmm_symint(self):
        from torch._subclasses.fake_tensor import FakeTensorMode

        shape_env = ShapeEnv()
        fake_mode = FakeTensorMode(shape_env=shape_env)

        B, M, K, N = [shape_env.create_unbacked_symint() for _ in range(4)]

        with fake_mode:
            A = torch.empty((B, M, K), device="meta")
            Bmat = torch.empty((B, K, N), device="meta")
            bias3 = torch.empty((B, M, N), device="meta")

            _ = torch.baddbmm(bias3, A, Bmat)


@skipIfTorchDynamo(
    "Creating ShapeEnv fails for confusing reasons (also we never expect dynamo to see code like this)"
)
class TestSymNumberMagicMethods(TestCase):
    def _do_test(self, fn, inp1, inp2, shape_env, is_unary_fn):
        with self.subTest(fn=fn, inp1=inp1, inp2=inp2, is_unary_fn=is_unary_fn):
            return self._do_test2(fn, inp1, inp2, shape_env, is_unary_fn)

    def _do_test2(self, fn, inp1, inp2, shape_env, is_unary_fn):
        # Helper function
        # NB: don't use one as that will get specialized
        # TODO: We don't have to circuitously create the float, can just
        # create a symfloat directly
        seed_node = (create_symint(shape_env, 2) / 2.0).node
        bool_seed_node = (create_symint(shape_env, 2) == 2).node

        def get_sym_inp(inp):
            # NB: this must come before int
            if isinstance(inp, bool):
                return torch.SymBool(to_node(bool_seed_node, inp))
            elif isinstance(inp, int):
                return torch.SymInt(to_node(seed_node, inp))
            else:
                return torch.SymFloat(to_node(seed_node, inp))

        if fn == "float_pow":
            if inp1 < 0:
                return

        if fn == "pow_by_natural":
            if isinstance(inp1, float) or isinstance(inp2, float):
                return
            if inp2 < 0:
                return

        def maybe_xfail(inp1, inp2):
            if fn == "sym_sqrt" and inp1 < 0:
                # ValueError: math domain error
                return self.assertRaises((ValueError,))
            elif (
                fn in ("float_truediv", "int_truediv", "int_floordiv", "mod")
                and inp2 == 0
            ):
                # ZeroDivisionError: division by zero
                return self.assertRaises((ZeroDivisionError,))
            elif fn in ["float_pow", "pow_by_natural"] and inp1 == 0 and inp2 < 0:
                # ZeroDivisionError: 0.0 cannot be raised to a negative power
                return self.assertRaises((ZeroDivisionError,))
            elif (
                # TODO: dear catastrophe waitress,
                # this doesn't work
                fn in ["float_pow", "pow_by_natural"]
                and inp1 < 0
                and (
                    type(inp1) is (SymInt, SymFloat) or type(inp2) is (SymInt, SymFloat)
                )
                and (type(inp1) is (SymFloat, float) or type(inp2) is (SymFloat, float))
            ):
                # Complex result, which we do not support:
                # TypeError: Cannot convert complex to float
                return self.assertRaises((RuntimeError,))
            elif fn in ("lshift", "rshift") and not (
                isinstance(inp1, (SymInt, int)) and isinstance(inp2, (SymInt, int))
            ):
                # TypeError: unsupported operand type(s)
                return self.assertRaises((TypeError,))
            elif fn in ("lshift", "rshift") and inp2 < 0:
                # ValueError: math domain error
                return self.assertRaises((ValueError,))
            else:
                return contextlib.nullcontext()

        lambda_apply = method_to_operator(fn)

        def guard_fn(v):
            if type(v) in (SymBool, bool):
                return guard_bool(v)
            elif type(v) in (SymFloat, float):
                return guard_float(v)
            else:  # SymInt, int
                return guard_int(v)

        # Get reference result
        with maybe_xfail(inp1, inp2):
            if is_unary_fn:
                ref_out = lambda_apply(inp1)
            else:
                ref_out = lambda_apply(inp1, inp2)

        # Symified first arg
        sym_inp1 = get_sym_inp(inp1)
        with maybe_xfail(sym_inp1, inp2):
            if is_unary_fn:
                out = lambda_apply(sym_inp1)
            else:
                out = lambda_apply(sym_inp1, inp2)
            self.assertTrue(isinstance(out, (SymInt, SymFloat, SymBool)))
            out = guard_fn(out)
            self.assertEqual(out, ref_out)

        if is_unary_fn:
            return

        # Symified second arg
        sym_inp2 = get_sym_inp(inp2)
        with maybe_xfail(inp1, sym_inp2):
            out = lambda_apply(inp1, sym_inp2)
            self.assertTrue(isinstance(out, (SymInt, SymFloat, SymBool)))
            out = guard_fn(out)
            self.assertEqual(out, ref_out)

        # Symified both args
        with maybe_xfail(sym_inp1, sym_inp2):
            out = lambda_apply(sym_inp1, sym_inp2)
            self.assertTrue(isinstance(out, (SymInt, SymFloat, SymBool)))
            out = guard_fn(out)
            self.assertEqual(out, ref_out)

    @parametrize("fn", list(sym_node.magic_methods.keys()))
    def test_bool_method(self, fn):
        # sym_ite has its own tests
        if fn not in sym_node.bool_magic_methods or fn == "sym_ite":
            self.skipTest(f"{fn} is non-bool")

        is_unary_fn = fn in sym_node.unary_methods
        shape_env = ShapeEnv()
        self._do_test(fn, True, False, shape_env, is_unary_fn)

    @parametrize("fn", list(sym_node.magic_methods.keys()))
    @parametrize("first_type", ["int", "float"])
    @parametrize("second_type", ["int", "float"])
    def test_method(self, fn, first_type, second_type):
        if first_type == "float":
            # TODO: Hmm, this looks like we skip all floats
            self.skipTest(f"{fn} is not a float magic method")

        if (
            first_type == "int" or second_type == "int"
        ) and fn in sym_node.only_float_magic_methods:
            self.skipTest(f"{fn} is not an int method")

        if second_type == "float" and fn == "mod":
            self.skipTest(f"{fn} only handles int")

        if fn in sym_node.bitwise_ops and (first_type != "int" or second_type != "int"):
            self.skipTest(f"{fn} is a bitwise op, only handles int")

        is_unary_fn = fn in sym_node.unary_methods or fn == "round"
        # Second argument is ignored for unary function. So only run for one type
        if is_unary_fn and second_type == "float":
            self.skipTest(f"{fn} is unary and already tested")

        if fn in sym_node.bool_magic_methods:
            self.skipTest(f"{fn} is bool")

        # Only floats here since these will be converted to int if necessary.
        # We also ignore complex and bool.
        values = (
            0.0,
            1.0,
            0.5 if fn in ("sym_acos", "sym_asin") else 2.5,  # avoid math domain error
        )

        neg_values = tuple(-x for x in values)

        for inp1, inp2 in itertools.chain(
            itertools.product(values, values),
            itertools.product(values, neg_values),
            itertools.product(neg_values, values),
            itertools.product(neg_values, neg_values),
        ):
            if first_type == "int":
                inp1 = int(inp1)
            if second_type == "int":
                inp2 = int(inp2)

            shape_env = ShapeEnv()

            self._do_test(fn, inp1, inp2, shape_env, is_unary_fn)

    def get_constant_bool(self, val):
        return SymBool(torch._C._get_constant_bool_symnode(val))

    @unittest.expectedFailure
    def test_symint_hashing(self):
        shape_env = ShapeEnv()
        hash(create_symint(shape_env, 3))

    def test_symnode_hashing(self):
        shape_env = ShapeEnv()

        # These all trigger specialization when hashed
        hash(create_symbool(shape_env, True))
        # We should be passing in float here, but create_symbol currently
        # only supports int
        hash(create_symfloat(shape_env, 3.0))

        # NestedInt (SymInt), constant SymBool, SymNode are hashable
        j1 = torch._C._get_nested_int(1, 1)
        j1_copy = torch._C._get_nested_int(1, 1)
        j2 = torch._C._get_nested_int(2, 1)
        t = self.get_constant_bool(True)
        t_copy = self.get_constant_bool(True)
        f = self.get_constant_bool(False)
        n = create_symint(shape_env, 3).node
        m = self.get_constant_bool(True).node

        self.assertIs(j1 == j1_copy, True)
        self.assertEqual(hash(j1), hash(j1_copy))
        self.assertIs(j1 == j2, False)
        self.assertNotEqual(hash(j1), hash(j2))
        self.assertIs(t == t_copy, True)
        self.assertEqual(hash(t), hash(t_copy))
        self.assertIs(t == f, False)
        self.assertNotEqual(hash(t), hash(f))

        hash(n)
        hash(m)

    def test_symint_deepcopy(self):
        shape_env = ShapeEnv()

        symnodes = (torch._C._get_nested_int(1, 1),)
        deepcopied_symnodes = copy.deepcopy(symnodes)
        self.assertEqual(symnodes, deepcopied_symnodes)

    def test_non_symbolic_symnode(self):
        j1 = torch._C._get_nested_int(1, 1)
        j2 = torch._C._get_nested_int(1, 1)
        j3 = torch._C._get_nested_int(3, 1)

        self.assertIsInstance(j1, torch.SymInt)
        self.assertNotIsInstance(j1, int)

        with self.assertRaisesRegex(
            RuntimeError, "add not supported by NestedIntSymNode"
        ):
            j1 + 3

        self.assertFalse(j1 == 3)
        with self.assertRaisesRegex(RuntimeError, "indeterminate"):
            self.assertFalse(3 >= j2)

        self.assertIs(j1 == j1, True)
        self.assertIs(j1 == j2, True)
        self.assertIs(j1 == j3, False)
        self.assertIs(j1 != j3, True)
        self.assertIs(j1 != j2, False)

        x = self.get_constant_bool(True)
        #
        # Unary
        #
        # op(constant SymBool)
        self.assertIs(x.__sym_not__(), False)

        #
        # Binary
        #
        # op(constant SymBool, bool)
        # op(constant SymBool, constant SymBool)
        # op(bool, constant SymBool)
        self.assertIs(operator.and_(x, True), True)
        self.assertIs(operator.and_(x, x), True)
        self.assertIs(operator.and_(True, x), True)

        # op(symbolic SymBool, constant Symbool)
        # op(constant SymBool, symbolic Symbool)
        shape_env = ShapeEnv()
        a = create_symint(shape_env, 2)
        b = create_symint(shape_env, 2)
        c = a == b  # symbolic SymBool
        d = self.get_constant_bool(True)
        e = operator.and_(c, d)
        f = operator.and_(d, c)
        self.assertTrue(is_symbolic(e))
        self.assertTrue(is_symbolic(f))
        self.assertIs(e.node.guard_bool("", 0), True)
        self.assertIs(f.node.guard_bool("", 0), True)

        # Comparing sizes
        sz1 = torch.Size([j1, j1, j1])
        sz2 = torch.Size([j1, j1, j1])
        self.assertIs(sz1 == sz2, True)

        sz1 = torch.Size([3, j1, 4])
        sz2 = torch.Size([3, j2, 4])
        self.assertIs(sz1 == sz2, True)
        self.assertIs(sz1 != sz2, False)

    def test_stride_symnode(self):
        shape_env = ShapeEnv()

        # check everything static
        t = create_fake_tensor_with_dynamic_size(
            torch.ones(3, 6),
            shape_env,
            dynamic_sizes=[
                DimDynamic.STATIC,
                DimDynamic.STATIC,
            ],
            dynamic_strides=[
                DimDynamic.INFER_STRIDE,
                DimDynamic.INFER_STRIDE,
            ],
        )
        self.assertTrue(all(isinstance(size, int) for size in t.size()))
        self.assertTrue(all(isinstance(stride, int) for stride in t.stride()))

        # check dynamic size but static dims
        t = create_fake_tensor_with_dynamic_size(
            torch.ones(3, 6),
            shape_env,
            dynamic_sizes=[
                DimDynamic.DYNAMIC,
                DimDynamic.DYNAMIC,
            ],
            dynamic_strides=[
                DimDynamic.INFER_STRIDE,
                DimDynamic.INFER_STRIDE,
            ],
        )
        # Expect stride to be inferred
        s0, s1 = t.size()
        s2, s3 = t.stride()
        self.assertTrue(isinstance(s0, torch.SymInt))
        self.assertTrue(isinstance(s1, torch.SymInt))
        self.assertTrue(isinstance(s2, torch.SymInt))
        self.assertTrue(s1 == s2)
        self.assertEqual(s3, 1)

        # Check dynamic stride but static dims
        t = create_fake_tensor_with_dynamic_size(
            torch.ones(3, 6),
            shape_env,
            dynamic_sizes=[
                DimDynamic.STATIC,
                DimDynamic.STATIC,
            ],
            dynamic_strides=[
                DimDynamic.DYNAMIC,
                DimDynamic.INFER_STRIDE,
            ],
        )
        s0, s1 = t.size()
        s2, s3 = t.stride()
        self.assertTrue(isinstance(s0, int))
        self.assertTrue(isinstance(s1, int))
        self.assertTrue(isinstance(s2, torch.SymInt))
        self.assertTrue(isinstance(s3, int))

        # Check dynamic sizes and dims, and ensure different symbol
        t = create_fake_tensor_with_dynamic_size(
            torch.ones(3, 6),
            shape_env,
            dynamic_sizes=[
                DimDynamic.DYNAMIC,
                DimDynamic.DYNAMIC,
            ],
            dynamic_strides=[
                DimDynamic.DYNAMIC,
                DimDynamic.INFER_STRIDE,
            ],
        )
        s0, s1 = t.size()
        s2, s3 = t.stride()
        self.assertTrue(isinstance(s0, torch.SymInt))
        self.assertTrue(isinstance(s1, torch.SymInt))
        self.assertTrue(isinstance(s2, torch.SymInt))
        self.assertTrue(isinstance(s3, int))
        self.assertTrue(str(s1.node.expr) != str(s2.node.expr))

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @parametrize("backend", ["inductor", "eager"])
    def test_dynamic_int_basic_compile(self, backend):
        from torch.fx.experimental.sym_node import DynamicInt

        cnt = CompileCounterWithBackend(backend)

        # test scalar inputs to function
        def f(x, y, z):
            out = torch.tensor([x + y + z])
            out = out + torch.zeros(abs(x) + 2).sum()  # test out tensor construction
            return out

        fn = torch.compile(f, fullgraph=True, backend=cnt)
        x = DynamicInt(1)
        z = DynamicInt(3)
        self.assertEqual(fn(x, x, z), f(1, 1, 3))  # guard: x == y
        self.assertEqual(fn(2, 2, 0), f(2, 2, 0))
        self.assertEqual(fn(-1, -1, 2), f(-1, -1, 2))
        self.assertEqual(cnt.frame_count, 1)  # no recompiles

        self.assertEqual(fn(3, 4, 5), f(3, 4, 5))  # now we recompile
        self.assertEqual(cnt.frame_count, 2)

        # test nn module property
        class Foo(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.i = DynamicInt(1)

            def forward(self, x):
                return torch.tensor([x + self.i])

        cnt.clear()
        m = Foo()
        mc = torch.compile(m, backend=cnt, fullgraph=True)

        self.assertEqual(mc(DynamicInt(0)), m(0))
        mc.i = -2  # override attribute
        self.assertEqual(mc(-1), m(-1))
        self.assertEqual(cnt.frame_count, 1)

    def test_dynamic_int_eager_usage(self):
        from torch.fx.experimental.sym_node import DynamicInt

        w = DynamicInt(-1)
        x = DynamicInt(0)
        y = DynamicInt(1)
        z = DynamicInt(2)

        def check(l, r):
            self.assertTrue(isinstance(l, DynamicInt))
            self.assertEqual(l, r)

        # test arithmetic
        check(2 * y + z, 4)
        check((10 - z) // 2, 4)
        check(1 // z, 0)
        check(-w + w**2, 2)
        check(x % z, 0)
        check(1 << z, 4)
        check(z | y, 3)
        check(min(y, z), 1)
        self.assertTrue(z > -2)
        with self.assertRaises(ZeroDivisionError):
            y % x

        # math, numpy
        self.assertEqual(math.cos(x), y)
        self.assertEqual(math.prod([z, z], start=z), 8)
        self.assertEqual(np.arange(z)[y], 1)
        self.assertTrue(np.allclose(np.ones([y, z]).sum(axis=x), np.ones(z)))

        # test conversions
        self.assertTrue(isinstance(x + 2, int))
        self.assertTrue(isinstance(x + 2, DynamicInt))
        self.assertEqual(y / 2.0, 0.5)  # this could return DynamicFloat in future
        self.assertEqual(float(z), 2.0)
        self.assertFalse(bool(x))
        self.assertEqual(DynamicInt(x).real, x.real)

        # torch functions, scalar inputs
        self.assertEqual(torch.arange(z)[:w][x], 0)
        self.assertEqual(torch.add(torch.tensor(w), torch.tensor(w), alpha=z), -3)
        self.assertEqual(
            list(torch.nn.Linear(z, y)(torch.randn(z * 2, z)).shape), [4, 1]
        )
        self.assertEqual(z * torch.ones(z).sum(dim=x), 4)


instantiate_parametrized_tests(TestSymNumberMagicMethods)


class TestFloorDiv(TestCase):
    @staticmethod
    def python_floordiv(x, y):
        return x // y

    @staticmethod
    def torch_floordiv(x, y):
        # Note: we fully evaluate here since FloorDiv might not always do
        # that.
        shape_env = ShapeEnv()
        return shape_env.evaluate_expr(FloorDiv(x, y))

    @staticmethod
    def yield_test_cases(values, negate=True):
        for x, y in values:
            yield (x, y)
            if negate:
                yield (-x, y)
                yield (x, -y)
                yield (-x, -y)

    def test_floordiv_float_int(self):
        values = ((7, 2),)

        for x, y in TestFloorDiv.yield_test_cases(values):
            self.assertEqual(
                TestFloorDiv.python_floordiv(x, y), TestFloorDiv.torch_floordiv(x, y)
            )

    def test_floordiv_div_by_one(self):
        values = ((2, 1),)

        for x, y in TestFloorDiv.yield_test_cases(values):
            self.assertEqual(
                TestFloorDiv.python_floordiv(x, y), TestFloorDiv.torch_floordiv(x, y)
            )

    def test_floordiv_div_does_not_generate_non_int_rational(self):
        s14 = sympy.Symbol("s14", integer=True, positive=True)
        s37 = sympy.Symbol("s37", integer=True, positive=True)

        inner_expr = FloorDiv(s14, 2016)
        middle_expr = (24 * s37 + 672) * inner_expr
        numerator = middle_expr + 21
        denominator = 22
        result = FloorDiv(numerator, denominator)
        rationals = result.atoms(sympy.Rational)
        all_rationals_ints = all(r.q == 1 for r in rationals)
        self.assertTrue(all_rationals_ints)

    def test_floordiv_simplify(self):
        # Tests how we simplify or evaluate FloorDiv without free variables
        shape_env = ShapeEnv()
        result = 21
        exprs = (7 * FloorDiv(6, 2),)

        for expr in exprs:
            self.assertEqual(expr, result)
            self.assertEqual(expr.doit(deep=False), result)
            self.assertEqual(expr.doit(deep=True), result)
            self.assertEqual(sympy.simplify(expr), result)
            self.assertEqual(shape_env.simplify(expr), result)
            self.assertEqual(shape_env.evaluate_expr(expr), result)

    def test_floordiv_assumptions(self):
        cases = (
            sympy.Symbol("i1", integer=True),
            sympy.Symbol("i2", integer=True),
        )

        for base, divisor in itertools.product(cases, repeat=2):

            def op():
                return FloorDiv(base, divisor)

            def is_complex(x):
                return x.is_integer is False and x.is_real is False and x.is_complex

            if is_complex(base) or is_complex(divisor):
                self.assertRaisesRegex(
                    TypeError,
                    (
                        r"unsupported operand type\(s\) for //: 'Symbol' and 'Symbol',"
                        r" expected integer or real"
                    ),
                    op,
                )
                continue

            op = op()

            # In regular Python, x//x == 1.0 if x is a float, but FloorDiv
            # always returns an integer 1 when both args are the same object.
            # This even works for Symbols with no assumptions specified.
            if base is divisor:
                self.assertTrue(op.is_integer)
                self.assertTrue(op.is_real)
            elif base.is_integer and divisor.is_integer:
                self.assertTrue(op.is_integer)
                self.assertTrue(op.is_real)
            else:
                self.assertEqual(op.is_integer, None)
                self.assertTrue(op.is_real)


class TestDimConstraints(TestCase):
    @skipIfTorchDynamo("mark_dynamic not supported")
    def test_simplify_max_1_0(self):
        x = torch.rand(10)
        torch._dynamo.mark_dynamic(x, 0, max=20, min=5)

        @torch.compile(fullgraph=True)
        def func(x, v):
            # test that statically_known_true
            if (v == 0 or v == 1) and not statically_known_true(
                max(v, (-1 + x.size()[0] // 2)) == (-1 + x.size()[0] // 2)
            ):
                raise AssertionError("error")

            if max(v, (-1 + x.size()[0] // 2)) == (-1 + x.size()[0] // 2):
                return x * 400
            else:
                return (x * 10) * 100

        # testing that this does not throw constraint violation error.
        self.assertEqual(func(x, 1), x * 400)
        self.assertEqual(func(x, 0), x * 400)

    def test_dim_constraints_reduce_congruences_simple(self):
        from sympy import Symbol

        s = Symbol("s", positive=True, integer=True)
        dim_constraints = DimConstraints({}, {}, set(), {})
        dim_constraints._congruences[s] = {
            (s / 2) % 2,
            (s / 2) % 8,
            (s / 2) % 4,
            s % 2,
            ((s / 16) + 2) % 4,
        }
        congruences = dim_constraints._reduce_congruences()
        self.assertEqual(congruences[s], {(s + 32) % 64})

    def test_dim_constraints_reduce_inequalities_simple(self):
        from sympy import Eq, Interval, Ne, Symbol
        from sympy.solvers.inequalities import reduce_inequalities

        s = Symbol("s", positive=True, integer=True)
        exprs = {
            s >= 2,
            Ne(8 * s, 16),
            Ne(s / 2, 1),
            Ne(16 * s, 32),
            s < 16,
            Ne(s, 2),
            s / 2 < 16,
            s / 2 > 1,
            s / 2 >= 2,
            Ne(3 * s / 2, 3),
        }
        solution = reduce_inequalities(exprs, s).as_set()
        self.assertEqual(solution, Interval.Ropen(4, 16))

        exprs.add(Eq(s / 2, 4))
        solution = reduce_inequalities(exprs, s).as_set()
        self.assertEqual(solution, {8})

    def test_dim_constraints_reduce_inequalities_error(self):
        from collections import defaultdict

        from sympy import Symbol
        from sympy.solvers.inequalities import reduce_inequalities

        from torch._dynamo.source import (
            LocalSource,
            TensorProperty,
            TensorPropertySource,
        )
        from torch.fx.experimental.symbolic_shapes import DynamicDimConstraintPrinter

        s0 = Symbol("s0", positive=True, integer=True)
        exprs = {
            4 * s0**3 - 4 * s0**2 + s0 <= 2147483647,
            s0 >= 2,
            s0**3 <= 2147483647,
            s0 <= 2147483647,
        }
        answer = reduce_inequalities(exprs, s0)

        symbol_to_source = defaultdict(list)
        symbol_to_source[s0].append(
            TensorPropertySource(
                base=LocalSource(local_name="a"), prop=TensorProperty.SIZE, idx=0
            )
        )
        dcp = DynamicDimConstraintPrinter(symbol_to_source, {})
        with self.assertRaisesRegex(
            AssertionError,
            "Unknown symbol.*created by constraints solver",
        ):
            dcp.doprint(answer)

    def test_dim_constraints_solve_full(self):
        from sympy import Eq, Integer, Ne, Symbol

        from torch._dynamo.source import (
            LocalSource,
            TensorProperty,
            TensorPropertySource,
        )

        src0 = TensorPropertySource(
            base=LocalSource(local_name="a"), prop=TensorProperty.SIZE, idx=0
        )
        src2 = TensorPropertySource(
            base=LocalSource(local_name="b"), prop=TensorProperty.SIZE, idx=0
        )
        src3 = TensorPropertySource(
            base=LocalSource(local_name="c"), prop=TensorProperty.SIZE, idx=0
        )
        src4 = TensorPropertySource(
            base=LocalSource(local_name="d"), prop=TensorProperty.SIZE, idx=0
        )

        src1 = TensorPropertySource(
            base=LocalSource(local_name="a"), prop=TensorProperty.SIZE, idx=2
        )
        src7 = TensorPropertySource(
            base=LocalSource(local_name="a"), prop=TensorProperty.SIZE, idx=3
        )

        src5 = TensorPropertySource(
            base=LocalSource(local_name="a"), prop=TensorProperty.SIZE, idx=1
        )
        src8 = TensorPropertySource(
            base=LocalSource(local_name="b"), prop=TensorProperty.SIZE, idx=1
        )

        src6 = TensorPropertySource(
            base=LocalSource(local_name="c"), prop=TensorProperty.SIZE, idx=1
        )
        src9 = TensorPropertySource(
            base=LocalSource(local_name="d"), prop=TensorProperty.SIZE, idx=1
        )
        src10 = TensorPropertySource(
            base=LocalSource(local_name="e"), prop=TensorProperty.SIZE, idx=1
        )

        src11 = TensorPropertySource(
            base=LocalSource(local_name="f"), prop=TensorProperty.SIZE, idx=1
        )
        src12 = TensorPropertySource(
            base=LocalSource(local_name="b"), prop=TensorProperty.SIZE, idx=2
        )

        s0 = Symbol("s0", positive=True, integer=True)
        s1 = Symbol("s1", positive=True, integer=True)
        s5 = Symbol("s5", positive=True, integer=True)
        s6 = Symbol("s6", positive=True, integer=True)
        symbol_to_source = {
            s0: [src0, src2, src3, src4],
            s1: [src1, src7],
            s5: [src5, src8],
            s6: [src6, src9, src10],
        }
        var_to_val = {s0: 8, s1: 96, s5: 22, s6: 21}
        marked_dynamic = {s0, s1, s5, s6}
        dim_constraints = DimConstraints(
            symbol_to_source, var_to_val, marked_dynamic, {}
        )
        dim_constraints.add_equality(src2, s0)
        dim_constraints.add_equality(src3, s0)
        dim_constraints.add_equality(src4, s0)
        dim_constraints.add_equality(src7, s1)
        dim_constraints.add_equality(src8, s5)
        dim_constraints.add_equality(src9, s6)
        dim_constraints.add_equality(src10, s6)
        dim_constraints.add_equality(src11, Integer(1))
        dim_constraints.add_equality(src12, Integer(3))

        dim_constraints.add(s1**2 <= 2147483647)
        dim_constraints.add(32 * s1**2 <= 2147483647)
        dim_constraints.add(s0 < 16)
        dim_constraints.add(Eq(Mod(s1, 2), 0))
        dim_constraints.add(Ne(FloorDiv(s1, 2), 1))
        dim_constraints.add(Ne((FloorDiv(s1, 2)) ** 2, 1))
        dim_constraints.add(32 * (FloorDiv(s1, 2)) ** 2 <= 2147483647)
        dim_constraints.add((FloorDiv(s1, 2)) ** 2 > 1)
        dim_constraints.add(Ne(FloorDiv(s1, 2), 1))
        dim_constraints.add(
            64 * (FloorDiv((FloorDiv(s1, 2) - 1), 2)) ** 2
            + 128 * (FloorDiv((FloorDiv(s1, 2) - 1), 2))
            + 64
            <= 2147483647
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 2) + 1, 1))
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 2)) ** 2
                + 2 * (FloorDiv((FloorDiv(s1, 2) - 1), 2))
                + 1,
                1,
            )
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 2) + 1, 1))
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 2)) ** 2
            + 2 * (FloorDiv((FloorDiv(s1, 2) - 1), 2))
            + 1
            > 1
        )
        dim_constraints.add(
            128 * (FloorDiv((FloorDiv(s1, 2) - 1), 4)) ** 2
            + 256 * (FloorDiv((FloorDiv(s1, 2) - 1), 4))
            + 128
            <= 2147483647
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 4) + 1, 1))
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 4)) ** 2
                + 2 * (FloorDiv((FloorDiv(s1, 2) - 1), 4))
                + 1,
                1,
            )
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 4) + 1, 1))
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 4)) ** 2
            + 2 * (FloorDiv((FloorDiv(s1, 2) - 1), 4))
            + 1
            > 1
        )
        dim_constraints.add(
            256 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            + 512 * (FloorDiv((FloorDiv(s1, 2) - 1), 8))
            + 256
            <= 2147483647
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) + 1, 1))
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                + 2 * (FloorDiv((FloorDiv(s1, 2) - 1), 8))
                + 1,
                1,
            )
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) + 1, 1))
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            + 2 * (FloorDiv((FloorDiv(s1, 2) - 1), 8))
            + 1
            > 1
        )
        dim_constraints.add(FloorDiv((FloorDiv(s1, 2) - 1), 8) + 1 >= 3)
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
            <= 2147483647
        )
        dim_constraints.add(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1 >= 0)
        dim_constraints.add(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1 >= 1)
        dim_constraints.add(
            Ne(
                60 * s0 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120 * s0 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60 * s0,
                0,
            )
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1, 1))
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 1,
                1,
            )
        )
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 1,
                0,
            )
        )
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 1
            >= 0
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1, 0))
        dim_constraints.add(
            1
            < 60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1, -1))
        dim_constraints.add(
            Ne(
                60 * s0 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120 * s0 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60 * s0,
                120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 120,
            )
        )
        dim_constraints.add(
            120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 120
            > 0
        )
        dim_constraints.add(
            Eq(
                60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2 * (Mod(s0, 2))
                - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8) * Mod(s0, 2)
                + 60 * (Mod(s0, 2)),
                0,
            )
        )
        dim_constraints.add(
            Ne(
                120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 120,
                0,
            )
        )
        dim_constraints.add(
            Ne(
                60
                * (FloorDiv(s0, 2))
                * (FloorDiv(s0, (FloorDiv(s0, 2))))
                * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120
                * FloorDiv(s0, 2)
                * FloorDiv(s0, (FloorDiv(s0, 2)))
                * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60 * (FloorDiv(s0, 2)) * (FloorDiv(s0, (FloorDiv(s0, 2)))),
                0,
            )
        )
        dim_constraints.add(Ne(FloorDiv(s0, 2), 1))
        dim_constraints.add(
            Ne(
                60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60,
                0,
            )
        )
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
            >= 0
        )
        dim_constraints.add(
            1
            < 60
            * (FloorDiv(s0, (FloorDiv(s0, 2))))
            * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv(s0, (FloorDiv(s0, 2))) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60 * (FloorDiv(s0, (FloorDiv(s0, 2))))
        )
        dim_constraints.add(Ne(16 * s0, 32))
        dim_constraints.add(Eq(16 * (Mod(s0, 2)), 0))
        dim_constraints.add(Ne(16 * s0, 32))
        dim_constraints.add(Eq(16 * (Mod(s0, 2)), 0))
        dim_constraints.add(FloorDiv(s0, 2) >= 2)
        dim_constraints.add(Ne(FloorDiv(s0, 2), 1))
        dim_constraints.add(1 < FloorDiv(s0, 2))
        dim_constraints.add(Ne(s0, 2))
        dim_constraints.add(
            60
            * (FloorDiv(s0, (FloorDiv(s0, 2))))
            * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv(s0, (FloorDiv(s0, 2))) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60 * (FloorDiv(s0, (FloorDiv(s0, 2))))
            >= 60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
        )
        dim_constraints.add(
            60
            * (FloorDiv(s0, 2))
            * (FloorDiv(s0, (FloorDiv(s0, 2))))
            * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120
            * FloorDiv(s0, 2)
            * FloorDiv(s0, (FloorDiv(s0, 2)))
            * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60 * (FloorDiv(s0, 2)) * (FloorDiv(s0, (FloorDiv(s0, 2))))
            > 0
        )
        dim_constraints.add(
            Ne(
                60
                * (FloorDiv(s0, 2))
                * (FloorDiv(s0, (FloorDiv(s0, 2))))
                * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120
                * FloorDiv(s0, 2)
                * FloorDiv(s0, (FloorDiv(s0, 2)))
                * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60 * (FloorDiv(s0, 2)) * (FloorDiv(s0, (FloorDiv(s0, 2)))),
                3 * (FloorDiv(s0, 2)) * (FloorDiv(s0, (FloorDiv(s0, 2)))),
            )
        )
        dim_constraints.add(
            Ne(
                20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 20,
                0,
            )
        )
        dim_constraints.add(
            20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 20
            >= 0
        )
        dim_constraints.add(
            Ne(
                20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 20,
                20,
            )
        )
        dim_constraints.add(
            Ne(
                20
                * (
                    Mod(
                        1,
                        (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                        - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                        + 1,
                    )
                ),
                0,
            )
        )
        dim_constraints.add(
            Ne(
                20
                * (FloorDiv((FloorDiv(s1, 2) - 1), 8))
                * (
                    Mod(
                        1,
                        (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                        / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                        - 2
                        * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                        / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                        + 1 / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1),
                    )
                )
                - 20
                * Mod(
                    1,
                    (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                    / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                    - 2
                    * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                    / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                    + 1 / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1),
                ),
                0,
            )
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1, 1))
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 1
            >= 1
        )
        dim_constraints.add(
            20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 20
            >= 0
        )
        dim_constraints.add(
            20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 20
            >= 1
        )
        dim_constraints.add(
            20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 20
            >= 2
        )
        dim_constraints.add(
            20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 20
            > 1
        )
        dim_constraints.add(
            20 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 40 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 20
            < 60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
        )
        dim_constraints.add(
            Ne(
                60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60,
                60,
            )
        )
        dim_constraints.add(
            Ne(
                FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1,
                (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 1,
            )
        )
        dim_constraints.add(
            Eq(
                (FloorDiv((FloorDiv(s1, 2) - 1), 8))
                * (
                    Mod(
                        (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                        / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                        - 2
                        * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                        / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                        + 1 / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1),
                        1,
                    )
                )
                - Mod(
                    (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                    / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                    - 2
                    * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                    / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1)
                    + 1 / (FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1),
                    1,
                ),
                0,
            )
        )
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 1,
                FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1,
            )
        )
        dim_constraints.add(Ne(8 * s0, 16))
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
            >= (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 1
        )
        dim_constraints.add(
            60
            * (FloorDiv(s0, (FloorDiv(s0, 2))))
            * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv(s0, (FloorDiv(s0, 2))) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60 * (FloorDiv(s0, (FloorDiv(s0, 2))))
            <= 2147483647
        )
        dim_constraints.add(
            90 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 180 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 90
            <= 2147483647
        )
        dim_constraints.add(FloorDiv(s0, 2) < 16)
        dim_constraints.add(FloorDiv(s0, 2) > 1)
        dim_constraints.add(
            Ne(
                90 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 180 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 90 * (FloorDiv(s0, 2)),
                0,
            )
        )
        dim_constraints.add(
            1
            < 90 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 180 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 90
        )
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 2 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 1
            > 1
        )
        dim_constraints.add(
            60
            * (FloorDiv(s0, (FloorDiv(s0, 2))))
            * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv(s0, (FloorDiv(s0, 2))) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60 * (FloorDiv(s0, (FloorDiv(s0, 2))))
            > 1
        )
        dim_constraints.add(
            Ne(
                60 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60 * (FloorDiv(s0, 2)),
                0,
            )
        )
        dim_constraints.add(
            90 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 180 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 90
            > 1
        )
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
            > 1
        )
        dim_constraints.add(
            Ne(
                60 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 120 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 60 * (FloorDiv(s0, 2)),
                3 * (FloorDiv(s0, 2)),
            )
        )
        dim_constraints.add(
            60 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60 * (FloorDiv(s0, 2))
            > 0
        )
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
            > 0
        )
        dim_constraints.add(
            Ne(
                120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 120,
                0,
            )
        )
        dim_constraints.add(
            1
            < 120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 120
        )
        dim_constraints.add(
            Ne(
                120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 120,
                6,
            )
        )
        dim_constraints.add(
            120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 120
            > 0
        )
        dim_constraints.add(
            Ne(
                120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 120,
                0,
            )
        )
        dim_constraints.add(
            120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 120
            <= 2147483647
        )
        dim_constraints.add(
            120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 120
            <= 20480
        )
        dim_constraints.add(
            Ne(
                90 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 180 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 90,
                0,
            )
        )
        dim_constraints.add(
            120 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 240 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 120
            > 1
        )
        dim_constraints.add(
            90 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 180 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 90
            <= 20480
        )
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 120 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 60
            <= 20480
        )
        dim_constraints.add(
            Ne(
                240 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 480 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 240,
                0,
            )
        )
        dim_constraints.add(Eq(6 * s5, 132))
        dim_constraints.add(Eq(4, FloorDiv(s0, 2)))
        dim_constraints.add(Eq(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1, 4))
        dim_constraints.add(
            Ne(
                64 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 128 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 64 * (FloorDiv(s0, 2)),
                0,
            )
        )
        dim_constraints.add(
            1
            < 64 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 128 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 64
        )
        dim_constraints.add(
            64 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 128 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 64
            <= 2147483647
        )
        dim_constraints.add(
            64 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 128 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 64
            > 1
        )
        dim_constraints.add(
            62 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 124 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 62
            <= 2147483647
        )
        dim_constraints.add(
            Ne(
                62 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 124 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 62 * (FloorDiv(s0, 2)),
                0,
            )
        )
        dim_constraints.add(
            1
            < 62 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 124 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 62
        )
        dim_constraints.add(Ne(3 * (FloorDiv(s0, 2)), 3))
        dim_constraints.add(Ne(3 * (FloorDiv(s0, 2)), 3))
        dim_constraints.add(Eq(FloorDiv(s0, 2), 4))
        dim_constraints.add(Eq(4, FloorDiv(s0, 2)))
        dim_constraints.add(Eq(FloorDiv(s0, 2), 4))
        dim_constraints.add(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 1 >= 3)
        dim_constraints.add(
            64 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 384 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 576
            <= 2147483647
        )
        dim_constraints.add(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 3 >= 0)
        dim_constraints.add(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 3 >= 1)
        dim_constraints.add(
            Ne(
                64 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 384 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 576 * (FloorDiv(s0, 2)),
                0,
            )
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 3, 1))
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 9,
                1,
            )
        )
        dim_constraints.add(
            Ne(
                (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 9,
                0,
            )
        )
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 9
            >= 0
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 3, 0))
        dim_constraints.add(
            1
            < 64 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 384 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 576
        )
        dim_constraints.add(Ne(FloorDiv((FloorDiv(s1, 2) - 1), 8) - 3, 1))
        dim_constraints.add(
            Ne(
                64 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 384 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 576 * (FloorDiv(s0, 2)),
                256,
            )
        )
        dim_constraints.add(
            Eq(
                64
                * (
                    Mod(
                        (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                        - 6 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                        + 9 * (FloorDiv(s0, 2)),
                        4,
                    )
                ),
                0,
            )
        )
        dim_constraints.add(
            Eq(
                FloorDiv(s0, 2),
                FloorDiv(
                    (
                        (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                        - 6 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                        + 9 * (FloorDiv(s0, 2))
                    ),
                    4,
                ),
            )
        )
        dim_constraints.add(
            Eq(
                FloorDiv(
                    (
                        (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                        - 6 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                        + 9 * (FloorDiv(s0, 2))
                    ),
                    4,
                ),
                FloorDiv(s0, 2),
            )
        )
        dim_constraints.add(
            Ne(64 * (Mod(FloorDiv((FloorDiv(s1, 2) - 1), 8) + 1, 4)), 0)
        )
        dim_constraints.add(
            Eq(
                64
                * (
                    Mod(
                        (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                        - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                        + 1,
                        4,
                    )
                ),
                0,
            )
        )
        dim_constraints.add(
            64 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 384 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 576 * (FloorDiv(s0, 2))
            > 0
        )
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 9
            >= 1
        )
        dim_constraints.add(
            Eq(
                64 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 384 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 576,
                256,
            )
        )
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 360 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 540
            <= 2147483647
        )
        dim_constraints.add(
            Ne(
                60 * (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 360 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 540 * (FloorDiv(s0, 2)),
                0,
            )
        )
        dim_constraints.add(
            1
            < 60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 360 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 540
        )
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 9
            <= 2147483647
        )
        dim_constraints.add(
            Ne(
                (FloorDiv(s0, 2)) * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
                - 6 * FloorDiv(s0, 2) * FloorDiv((FloorDiv(s1, 2) - 1), 8)
                + 9 * (FloorDiv(s0, 2)),
                0,
            )
        )
        dim_constraints.add(
            1
            < (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 9
        )
        dim_constraints.add(
            (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 6 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 9
            > 1
        )
        dim_constraints.add(
            60 * (FloorDiv((FloorDiv(s1, 2) - 1), 8)) ** 2
            - 360 * FloorDiv((FloorDiv(s1, 2) - 1), 8)
            + 540
            > 1
        )
        dim_constraints.add(s0 >= 2)
        dim_constraints.add(s1 >= 2)
        dim_constraints.add(s6 >= 2)
        dim_constraints.add(s5 >= 2)

        dim_constraints.solve()
        self.assertEqual(
            dim_constraints._static_results,
            {
                "L['c'].size()[0] == 8",
                "L['d'].size()[0] == 8",
                "L['a'].size()[2] == 96",
                "L['f'].size()[1] == 1",
                "L['a'].size()[3] == 96",
                "L['b'].size()[2] == 3",
                "L['b'].size()[1] == 22",
                "L['b'].size()[0] == 8",
                "L['a'].size()[1] == 22",
                "L['a'].size()[0] == 8",
            },
        )
        self.assertEqual(
            dim_constraints._dynamic_results,
            {
                "2 <= L['c'].size()[1]",
                "L['d'].size()[1] == L['c'].size()[1]",
                "L['e'].size()[1] == L['c'].size()[1]",
            },
        )


class TestGuardsExpressions(TestCase):
    """
    Tests the guards-related methods used by the inductor FX graph cache.
    """

    def test_guards_gt_lt(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 6)
        s1 = create_symint(shape_env, 7)
        s2 = create_symint(shape_env, 5)

        guard_int(sym_int(s0 > 5))
        guard_int(sym_int(s0 < 7))

        guards = shape_env.produce_guards_expression([s0])

        self.assertTrue(
            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s0)])
        )
        self.assertFalse(
            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s1)])
        )
        self.assertFalse(
            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s2)])
        )

    def test_guards_float_print(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 3)
        guard_bool(2 / s0 == 2 / 3)
        guards = shape_env.produce_guards_expression([s0])
        self.assertTrue(
            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s0)])
        )

    @skipIfTorchDynamo("Not a TorchDynamo suitable test")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_guard_or_true(self):
        from torch.fx.experimental.symbolic_shapes import guard_or_true

        def func(a, b):
            x = a.item()
            if guard_or_true(x == 1):
                return b * 10
            else:
                return b * 20

        # eager.
        self.assertEqual(func(torch.tensor([1]), torch.tensor([1])), torch.tensor([10]))
        self.assertEqual(func(torch.tensor([2]), torch.tensor([1])), torch.tensor([20]))

        # compile with unbacked.
        unbacked_func = torch.compile(func, dynamic=True, fullgraph=True)
        a = torch.tensor([1])
        b = torch.tensor([1])
        unbacked_func(a, b)

        # always return b*10
        self.assertEqual(
            unbacked_func(torch.tensor([1]), torch.tensor([1])), torch.tensor([10])
        )
        self.assertEqual(
            unbacked_func(torch.tensor([2]), torch.tensor([1])), torch.tensor([10])
        )

        # Test that statically known true works.
        def func2(a, b):
            x = a.item()
            if guard_or_true(x != x):
                return b * 10
            else:
                return b * 20

        unbacked_func2 = torch.compile(func2, dynamic=True, fullgraph=True)
        a = torch.tensor([1])
        b = torch.tensor([1])
        unbacked_func2(a, b)
        # always return b*20
        self.assertEqual(
            unbacked_func2(torch.tensor([1]), torch.tensor([1])), torch.tensor([20])
        )
        self.assertEqual(
            unbacked_func2(torch.tensor([2]), torch.tensor([1])), torch.tensor([20])
        )

        # Test backed_size_oblivious
        with torch.fx.experimental._config.patch("backed_size_oblivious", True):

            def func3(a, b):
                if guard_or_true(a.size()[0] != 9):
                    return b * 10
                else:
                    return b * 20

            compiled = torch.compile(func3, dynamic=True, fullgraph=True)
            a = torch.rand(9, 2)
            b = torch.rand(3, 4)

            self.assertEqual(func3(a, b), b * 20)
            self.assertEqual(compiled(a, b), b * 10)

    @skipIfTorchDynamo("Not a TorchDynamo suitable test")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_guard_or_false(self):
        from torch.fx.experimental.symbolic_shapes import guard_or_false

        def func(a, b):
            x = a.item()
            if guard_or_false(x == 1):
                return b * 10
            else:
                return b * 20

        # eager.
        self.assertEqual(func(torch.tensor([1]), torch.tensor([1])), torch.tensor([10]))
        self.assertEqual(func(torch.tensor([2]), torch.tensor([1])), torch.tensor([20]))

        # compile with unbacked.
        unbacked_func = torch.compile(func, dynamic=True, fullgraph=True)
        a = torch.tensor([1])
        b = torch.tensor([1])
        unbacked_func(a, b)

        # always return b*20
        self.assertEqual(
            unbacked_func(torch.tensor([1]), torch.tensor([1])), torch.tensor([20])
        )
        self.assertEqual(
            unbacked_func(torch.tensor([2]), torch.tensor([1])), torch.tensor([20])
        )

        # Test that statically known true works.
        def func2(a, b):
            x = a.item()
            if guard_or_false(x == x):
                return b * 10
            else:
                return b * 20

        unbacked_func2 = torch.compile(func2, dynamic=True, fullgraph=True)
        a = torch.tensor([1])
        b = torch.tensor([1])
        unbacked_func2(a, b)
        # always return b*10
        self.assertEqual(
            unbacked_func2(torch.tensor([1]), torch.tensor([1])), torch.tensor([10])
        )
        self.assertEqual(
            unbacked_func2(torch.tensor([2]), torch.tensor([1])), torch.tensor([10])
        )

        # Test backed_size_oblivious
        with torch.fx.experimental._config.patch("backed_size_oblivious", True):

            def func3(a, b):
                if guard_or_false(a.size()[0] == 9):
                    return b * 10
                else:
                    return b * 20

            compiled = torch.compile(func3, dynamic=True, fullgraph=True)
            a = torch.rand(9, 2)
            b = torch.rand(3, 4)

            self.assertEqual(func3(a, b), b * 10)
            self.assertEqual(compiled(a, b), b * 20)

    def test_guards_float_div(self):
        shape_env = ShapeEnv()
        s0 = create_symint(shape_env, 8)
        s1 = create_symint(shape_env, 7)

        guard_int(sym_int(s0 / 2.0))
        guards = shape_env.produce_guards_expression([s0])

        self.assertIn("math.trunc(", guards)
        self.assertIn("float(", guards)
        self.assertTrue(
            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s0)])
        )
        self.assertFalse(
            shape_env.evaluate_guards_expression(guards, [guarding_hint_or_throw(s1)])
        )

    @unittest.skipIf(
        TEST_XPU, "Skipped on XPU"
    )  # https://github.com/intel/torch-xpu-ops/issues/2169"
    @skipIfTorchDynamo("Attempt to trace generator")
    @torch.fx.experimental._config.patch("use_duck_shape", False)
    def test_size_comparison_no_recompile(self):
        """
        Test that size comparisons don't cause recompilation.
        When comparing x.size() == b.size() with different sizes,
        the compiled function should only compile once.
        We should not guard in sizes of the inner elements.
        """
        cnt = CompileCounter()

        @torch.compile(fullgraph=True, dynamic=True, backend=cnt)
        def f(x, b):
            if x.size() == b.size():
                return x
            return x * 2

        # First call: shapes differ (1, 2) vs (2, 4, 9), so if branch is False
        f(torch.rand(10, 2), torch.rand(20, 4, 9))

        # Second call: shapes differ again (1, 2) vs (1, 4, 9), so if branch is False
        f(torch.rand(10, 2), torch.rand(10, 4, 9))

        # Should only compile once despite different input shapes
        self.assertEqual(
            cnt.frame_count,
            1,
            f"Expected 1 compilation, got {cnt.frame_count}. "
            f"Size comparison should not cause recompilation.",
        )

    @skipIfTorchDynamo("Test uses torch.compile")
    def test_python_mod_padding_no_recompile(self):
        """
        Test that padding with PythonMod and slicing back to original size
        doesn't cause guards or recompilation.
        """
        cnt = CompileCounter()

        @torch.compile(backend=cnt, dynamic=True)
        def func(x):
            # Padding to align to multiple of 4
            padding = (-x.size()[0]) % 4
            aligned_size = x.size()[0] + padding
            tensor = torch.empty(aligned_size, x.size()[1])
            # Do some work on padded tensor
            tensor = tensor * 100
            # Remove padding with slicing
            remove_padding = tensor[0 : x.size()[0], ...]
            # View should work without guards since shapes match
            return remove_padding.view(x.size())

        # First call we pass something does not need padding
        result1 = func(torch.rand(4, 10))
        self.assertEqual(result1.shape, (4, 10))

        # Second call with different size needs padding
        result1 = func(torch.rand(10, 12))
        self.assertEqual(result1.shape, (10, 12))

        # Should only compile once despite different input shapes
        self.assertEqual(
            cnt.frame_count,
            1,
            f"Expected 1 compilation, got {cnt.frame_count}. "
            f"PythonMod padding should not cause recompilation.",
        )

        # test with unbacked.
        torch._dynamo.reset()

        a = torch.rand(10, 10)
        torch._dynamo.decorators.mark_unbacked(a, 0)
        torch._dynamo.decorators.mark_unbacked(a, 1)
        result1 = func(a)
        self.assertEqual(result1.shape, (10, 10))

        result1 = func(torch.rand(4, 10))
        self.assertEqual(result1.shape, (4, 10))

        self.assertEqual(cnt.frame_count, 2)

    def test_remove_symbols_without_guarding(self):
        from torch._functorch.partitioners import _remove_symbols_without_guarding

        shape_env = ShapeEnv()

        x = create_fake_tensor_with_dynamic_size(
            torch.randn(5, 8),
            shape_env,
            dynamic_sizes=[
                DimDynamic.DYNAMIC,
                DimDynamic.DYNAMIC,
            ],
            dynamic_strides=[
                DimDynamic.INFER_STRIDE,
                DimDynamic.INFER_STRIDE,
            ],
        )

        self.assertEqual(f"{x.stride()}", "(s49, 1)")
        self.assertEqual(f"{x.shape}", "torch.Size([s26, s49])")

        x_clean = _remove_symbols_without_guarding(x, 4096)

        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":
            if node.meta["val"].size()[0] != 2:
                raise AssertionError(
                    f"expected size()[0] == 2, got {node.meta['val'].size()[0]}"
                )
    return graph


class TestUnbacked(TestCase):
    @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/156135")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @parametrize("backend", ["inductor", "eager"])
    def test_deferred_neq_assert(self, backend):
        @torch.compile(fullgraph=True, backend=backend)
        def func(a):
            torch._check(a.item() != 5)
            return a.item() * 10

        func(torch.tensor([100]))

        with self.assertRaises(RuntimeError):
            func(torch.tensor([5]))

    # Test a situation where we generate a runtime assert i.e: u1==s1, then we specialize s1
    # later on to a constant.
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @parametrize("backend", ["inductor", "eager"])
    def test_post_specialize_runtime_assert1(self, backend):
        @torch.compile(dynamic=True, backend=backend)
        def func(x, y):
            u0 = y.item()
            s0 = x.size()[0]
            s1 = x.size()[1]
            torch._check(u0 + s0 + s1 == 102)
            if s0 != 2:
                raise AssertionError(f"expected s0 == 2, got {s0}")
            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)
    @parametrize("backend", ["inductor", "eager"])
    def test_post_specialize_runtime_assert2(self, backend):
        @torch.compile(dynamic=True, backend=backend)
        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]))

    @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/156135")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @parametrize("backend", ["inductor", "eager"])
    def test_deferred_sym_or_assert(self, backend):
        @torch.compile(fullgraph=True, backend=backend)
        def func(a, b):
            torch._check(operator.or_(a.item() == 5, b.item() == 5))
            return a.item() * 10

        func(torch.tensor([5]), torch.tensor([100]))
        func(torch.tensor([100]), torch.tensor([5]))

    def test_has_free_symbols(self):
        self.assertFalse(has_free_symbols(sympy.S.true))
        self.assertFalse(has_free_symbols(sympy.Max(1, 10, evaluate=False)))

        self.assertFalse(has_free_symbols(sympy.sympify("1")))
        self.assertFalse(has_free_symbols(sympy.sympify("1.1")))
        self.assertTrue(has_free_symbols(sympy.sympify("a")))
        self.assertTrue(has_free_symbols(sympy.sympify("a*2")))
        self.assertTrue(has_free_symbols(sympy.sympify("a+b")))

    @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/156135")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @parametrize("backend", ["inductor", "eager"])
    def test_deferred_sym_eq_assert(self, backend):
        @torch.compile(fullgraph=True, backend=backend)
        def func(a, b):
            torch._check(b.item() == 5)
            return a * 10

        func(torch.tensor([5]), torch.tensor([5]))
        with self.assertRaises(RuntimeError):
            func(torch.tensor([100]), torch.tensor([1]))

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @parametrize("backend", ["inductor", "eager"])
    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_deferred_with_unbacked_input(self, backend):
        @torch.compile(fullgraph=True, dynamic=True, backend=backend)
        def func(a, b):
            torch._check(a.size()[0] == b.size()[0])
            return a * 10

        a = torch.rand(1, 1)
        b = torch.rand(1, 1)
        torch._dynamo.decorators.mark_unbacked(a, 0)
        torch._dynamo.decorators.mark_unbacked(b, 0)
        func(a, b)

        # inductor adds the check sometimes itself so it will be reflected
        # as AssertionError.
        with self.assertRaises((AssertionError, RuntimeError)):
            func(a, torch.rand(2, 1))

    @pytest.mark.xfail(reason="https://github.com/pytorch/pytorch/issues/163785")
    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_do_not_guard_unbacked_inputs(self):
        @torch.compile(fullgraph=True, dynamic=True, backend="inductor")
        def func(a, b):
            a.expand(b.shape)
            return a * 10

        a = torch.rand(1, 1)
        b = torch.rand(1, 1)

        torch._dynamo.decorators.mark_unbacked(a, 0)
        torch._dynamo.decorators.mark_unbacked(a, 1)
        torch._dynamo.decorators.mark_unbacked(b, 0)
        torch._dynamo.decorators.mark_unbacked(b, 1)

        log_stream, ctx = logs_to_string("torch._dynamo.guards", "guards")
        with ctx():
            func(a, b)
            func(torch.rand(4, 5), torch.rand(4, 5))

        guards = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertFalse("SYMBOLIC_SHAPE_GUARD" in guards)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_div_unbacked_eq_input_tensors(self):
        @torch.compile(fullgraph=True)
        def func(a, b):
            x = a.size()[0]
            y = b.size()[0]
            torch._check(x == y)
            if x // y == 1:
                a = a * 10
            if 2 * x // y == 2:
                a = a * 20
            return a

        a = torch.randn(10, 10)
        b = torch.randn(10, 20)

        torch._dynamo.decorators.mark_unbacked(a, 0)
        torch._dynamo.decorators.mark_unbacked(b, 0)
        func(a, b)

    @torch.compiler.config.patch(unbacked_sources="L['x'],L['y']")
    def test_div_unbacked_eq_input_ints(self):
        @torch.compile(fullgraph=True)
        def func(x, y):
            a = torch.rand(1)
            torch._check(x == y)
            if x // y == 1:
                a = a * 10
            if 2 * x // y == 2:
                a = a * 20
            return a

        func(10, 10)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    @torch.compiler.config.patch(unbacked_sources="L['y']")
    def test_div_unbacked_eq_globals(self):
        tensor = torch.rand(10, 44)
        y = 10

        @torch.compile(fullgraph=True, dynamic=True)
        def func():
            a = torch.rand(1)
            x = tensor.size()[0]
            torch._check(x == y)
            if x // y == 1:
                a = a * 10
            if 2 * x // y == 2:
                a = a * 20
            return a

        torch._dynamo.decorators.mark_unbacked(tensor, 0)
        func()

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_div_unbacked_eq_item(self):
        @torch.compile(fullgraph=True)
        def func(a, b):
            x = a.item()
            y = b.item()
            torch._check(x == y)
            # TODO we should not need those torch checks.
            torch._check(x // y == 1)
            torch._check(2 * x // y == 2)
            if x // y == 1:
                a = a * 10
            if 2 * x // y == 2:
                a = a * 20
            return a

        a = torch.tensor([1])
        b = torch.tensor([1])
        func(a, b)

    def test_ignorable_fresh_unbacked_symbols(self):
        # Test that symbols created during fake tensor tracing but not returned
        # can be marked as ignorable to avoid unbound symbol errors

        # Positive case: marking symbol as ignorable should work
        with torch.library._scoped_library("test_ignorable", "FRAGMENT") as lib:
            torch.library.define(
                "test_ignorable::custom_op",
                "(Tensor x, bool y) -> Tensor",
                lib=lib,
            )

            @torch.library.register_fake("test_ignorable::custom_op", lib=lib)
            def custom_op_fake(x, y):
                ctx = torch._custom_op.impl.get_ctx()
                # Create a symint during fake tensor tracing
                if y:
                    with ctx._shape_env.ignore_fresh_unbacked_symbols():
                        sz = ctx.new_dynamic_size()
                else:
                    sz = ctx.new_dynamic_size()
                # Return something else (not the symint we created)
                return x.clone()

            @torch.library.impl(
                "test_ignorable::custom_op",
                "CompositeExplicitAutograd",
                lib=lib,
            )
            def custom_op_impl(x, y):
                return x.clone()

            @torch.compile(fullgraph=True)
            def func_positive(x):
                return torch.ops.test_ignorable.custom_op(x, True)

            @torch.compile(fullgraph=True)
            def func_negative(x):
                return torch.ops.test_ignorable.custom_op(x, False)

            x = torch.randn(3, 4)

            # This should not raise an error about unbound symbols
            result = func_positive(x)
            self.assertEqual(result, x)

            # This should raise an error about unbound symbols
            with self.assertRaisesRegex(
                Exception, "Pending unbacked symbols.*not in returned outputs"
            ):
                func_negative(x)

    def test_meta_copy(self):
        """
        Test that meta_copy_ does not raise when self and src have different
        unbacked symint sizes.
        """
        from torch._meta_registrations import meta_copy_
        from torch._subclasses.fake_tensor import FakeTensorMode

        shape_env = ShapeEnv()
        fake_mode = FakeTensorMode(shape_env=shape_env)
        with fake_mode:
            u0 = shape_env.create_unbacked_symint()
            u1 = shape_env.create_unbacked_symint()

            self_tensor = torch.empty((u0, 256))
            src_tensor = torch.empty((u1, 256))

            meta_copy_(self_tensor, src_tensor)


class TestUbackedOps(TestCase):
    @fresh_cache()
    @skipIfTorchDynamo("not allowed to trace mark_unbacked")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_reshape1(self):
        cnt = CompileCounterWithBackend("inductor")

        # Reshape happens in place reshape (no-clone)
        # reshape u1 -> (u0*u0)
        def func(x, y):
            f = y.item()
            t1 = x.view((f, f))
            t2 = x.reshape((f, f))
            t3 = torch._ops.ops.aten.view_copy(x, (f, f))
            return t1 * 10, t2 * 10, t3

        compiled_func = torch.compile(
            fullgraph=True,
            backend=cnt,
            dynamic=True,
        )(func)

        # create a non-contiguous with data being even numbers in [0:cnt-1]
        # and reshape it into sqrt(cnt)*sqrt(cnt)
        def make_non_contiguous_tensor_and_test(cnt):
            # create a non-contiguous tensor x that is skipping odd indices.
            x = torch.arange(cnt * 2)
            x = x.as_strided((x.size()[0] // 2,), (2,))

            torch._dynamo.decorators.mark_unbacked(x, 0)
            sz = torch.tensor([int(math.sqrt(cnt))])
            compiled_result = compiled_func(x, sz)
            eager_result = func(x, sz)
            self.assertEqual(compiled_result, eager_result)

        log_stream, ctx = logs_to_string(
            "torch._functorch._aot_autograd.graph_capture", "aot_graphs"
        )
        with ctx():
            make_non_contiguous_tensor_and_test(4)
        aot_graphs = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            aot_graphs,
            """\
def forward(self, arg0_1: "i64[1][1]cpu", arg1_1: "Sym(u1)", arg2_1: "Sym(s7)", arg3_1: "i64[u1][s7]cpu"):
        ge: "Sym(u1 >= 0)" = arg1_1 >= 0
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u1 >= 0 on node 'ge'");  ge = _assert_scalar = None
        _local_scalar_dense: "Sym(u0)" = torch.ops.aten._local_scalar_dense.default(arg0_1);  arg0_1 = None
        ge_1: "Sym(u0 >= 0)" = _local_scalar_dense >= 0
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u0 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
        pow_1: "Sym(u0**2)" = _local_scalar_dense ** 2
        eq: "Sym(Eq(u1, u0**2))" = arg1_1 == pow_1;  arg1_1 = pow_1 = None
        _assert_scalar_2 = torch.ops.aten._assert_scalar.default(eq, "Runtime assertion failed for expression Eq(u1, u0**2) on node 'eq'");  eq = _assert_scalar_2 = None
        view: "i64[u0, u0][s7*u0, s7]cpu" = torch.ops.aten.view.default(arg3_1, [_local_scalar_dense, _local_scalar_dense])
        view_1: "i64[u0, u0][s7*u0, s7]cpu" = torch.ops.aten.view.default(arg3_1, [_local_scalar_dense, _local_scalar_dense])
        view_2: "i64[u0, u0][s7*u0, s7]cpu" = torch.ops.aten.view.default(arg3_1, [_local_scalar_dense, _local_scalar_dense]);  arg3_1 = _local_scalar_dense = None
        clone: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.clone.default(view_2);  view_2 = None
        mul_11: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.mul.Tensor(view, 10);  view = None
        mul_14: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.mul.Tensor(view_1, 10);  view_1 = None
        return (mul_11, mul_14, clone)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

        make_non_contiguous_tensor_and_test(49)
        self.assertEqual(cnt.frame_count, 1)

        # Pass in a contiguous tensor, it will recompile due to stride being 1 (0/1 specialization).
        # marking strides unbacked would have avoided the recompilation here.
        x = torch.arange(100)
        torch._dynamo.decorators.mark_unbacked(x, 0)

        log_stream, ctx = logs_to_string(
            "torch._functorch._aot_autograd.graph_capture", "aot_graphs"
        )
        with ctx():
            compiled_result = compiled_func(x, torch.tensor([10]))
            eager_result = func(x, torch.tensor([10]))
            self.assertEqual(compiled_result, eager_result)
            self.assertEqual(cnt.frame_count, 2)

        aot_graphs = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            aot_graphs,
            """\
def forward(self, arg0_1: "i64[1][1]cpu", arg1_1: "Sym(u1)", arg2_1: "i64[u1][1]cpu"):
        ge: "Sym(u1 >= 0)" = arg1_1 >= 0
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u1 >= 0 on node 'ge'");  ge = _assert_scalar = None
        _local_scalar_dense: "Sym(u0)" = torch.ops.aten._local_scalar_dense.default(arg0_1);  arg0_1 = None
        ge_1: "Sym(u0 >= 0)" = _local_scalar_dense >= 0
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u0 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
        pow_1: "Sym(u0**2)" = _local_scalar_dense ** 2
        eq: "Sym(Eq(u1, u0**2))" = arg1_1 == pow_1;  arg1_1 = pow_1 = None
        _assert_scalar_2 = torch.ops.aten._assert_scalar.default(eq, "Runtime assertion failed for expression Eq(u1, u0**2) on node 'eq'");  eq = _assert_scalar_2 = None
        view: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.view.default(arg2_1, [_local_scalar_dense, _local_scalar_dense])
        view_1: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.view.default(arg2_1, [_local_scalar_dense, _local_scalar_dense])
        view_2: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.view.default(arg2_1, [_local_scalar_dense, _local_scalar_dense]);  arg2_1 = _local_scalar_dense = None
        clone: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.clone.default(view_2);  view_2 = None
        mul_6: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.mul.Tensor(view, 10);  view = None
        mul_9: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten.mul.Tensor(view_1, 10);  view_1 = None
        return (mul_6, mul_9, clone)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

        x = torch.arange(25)
        compiled_result = compiled_func(x, torch.tensor([5]))
        eager_result = func(x, torch.tensor([5]))
        self.assertEqual(cnt.frame_count, 2)

    @skipIfTorchDynamo("not allowed to trace mark_unbacked")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_reshape2(self):
        cnt = CompileCounterWithBackend("inductor")

        # This reshape requires a clone when the input is not contiguous and we can't compute strides.
        # reshape (u2, u3) -> (u0, u1)
        def func(x, y):
            u0, u1 = y.tolist()

            result1 = torch.reshape(x, (u0, u1))
            return result1 * 10

        compiled_func = torch.compile(fullgraph=True, backend=cnt, dynamic=True)(func)

        x = torch.randn(10, 10)
        # make x not contiguous.
        x = x.t_()
        torch._dynamo.decorators.mark_unbacked(x, 0)
        torch._dynamo.decorators.mark_unbacked(x, 1)

        log_stream, ctx = logs_to_string(
            "torch._functorch._aot_autograd.graph_capture", "aot_graphs"
        )
        with ctx():
            result_eager = func(x, torch.tensor([5, 20]))
            result_compiled = compiled_func(x, torch.tensor([5, 20]))
            self.assertEqual(result_compiled, result_eager)
            self.assertEqual(cnt.frame_count, 1)

        aot_graphs = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            aot_graphs,
            """\
def forward(self, arg0_1: "i64[2][1]cpu", arg1_1: "Sym(u2)", arg2_1: "Sym(u3)", arg3_1: "f32[u2, u3][1, u2]cpu"):
        ge: "Sym(u2 >= 0)" = arg1_1 >= 0
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u2 >= 0 on node 'ge'");  ge = _assert_scalar = None
        ge_1: "Sym(u3 >= 0)" = arg2_1 >= 0
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u3 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
        select: "i64[][]cpu" = torch.ops.aten.select.int(arg0_1, 0, 0)
        _local_scalar_dense: "Sym(u0)" = torch.ops.aten._local_scalar_dense.default(select);  select = None
        ge_2: "Sym(u0 >= 0)" = _local_scalar_dense >= 0
        _assert_scalar_2 = torch.ops.aten._assert_scalar.default(ge_2, "Runtime assertion failed for expression u0 >= 0 on node 'ge_2'");  ge_2 = _assert_scalar_2 = None
        sym_sum: "Sym(u0 + 1)" = torch.sym_sum((1, _local_scalar_dense))
        gt: "Sym(u0 + 1 > 0)" = sym_sum > 0;  sym_sum = None
        _assert_scalar_3 = torch.ops.aten._assert_scalar.default(gt, "Runtime assertion failed for expression 0 < u0 + 1 on node 'gt'");  gt = _assert_scalar_3 = None
        select_1: "i64[][]cpu" = torch.ops.aten.select.int(arg0_1, 0, 1);  arg0_1 = None
        _local_scalar_dense_1: "Sym(u1)" = torch.ops.aten._local_scalar_dense.default(select_1);  select_1 = None
        ge_3: "Sym(u1 >= 0)" = _local_scalar_dense_1 >= 0
        _assert_scalar_4 = torch.ops.aten._assert_scalar.default(ge_3, "Runtime assertion failed for expression u1 >= 0 on node 'ge_3'");  ge_3 = _assert_scalar_4 = None
        sym_sum_1: "Sym(u1 + 1)" = torch.sym_sum((1, _local_scalar_dense_1))
        gt_1: "Sym(u1 + 1 > 0)" = sym_sum_1 > 0;  sym_sum_1 = None
        _assert_scalar_5 = torch.ops.aten._assert_scalar.default(gt_1, "Runtime assertion failed for expression 0 < u1 + 1 on node 'gt_1'");  gt_1 = _assert_scalar_5 = None
        mul: "Sym(u2*u3)" = arg1_1 * arg2_1;  arg1_1 = arg2_1 = None
        mul_1: "Sym(u0*u1)" = _local_scalar_dense * _local_scalar_dense_1
        eq: "Sym(Eq(u2*u3, u0*u1))" = mul == mul_1;  mul = mul_1 = None
        _assert_scalar_6 = torch.ops.aten._assert_scalar.default(eq, "Runtime assertion failed for expression Eq(u2*u3, u0*u1) on node 'eq'");  eq = _assert_scalar_6 = None
        clone: "f32[u2, u3][Max(1, u3), 1]cpu" = torch.ops.aten.clone.default(arg3_1, memory_format = torch.contiguous_format);  arg3_1 = None
        view: "f32[u0, u1][Max(1, u1), 1]cpu" = torch.ops.aten.view.default(clone, [_local_scalar_dense, _local_scalar_dense_1]);  clone = _local_scalar_dense = _local_scalar_dense_1 = None
        mul_21: "f32[u0, u1][Max(1, u1), 1]cpu" = torch.ops.aten.mul.Tensor(view, 10);  view = None
        return (mul_21,)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

        result_eager = func(x, torch.tensor([2, 50]))
        result_compiled = compiled_func(x, torch.tensor([2, 50]))
        self.assertEqual(result_compiled, result_eager)
        self.assertEqual(cnt.frame_count, 1)

        x = torch.randn(4, 4).t_()
        result_eager = func(x, torch.tensor([2, 8]))
        result_compiled = compiled_func(x, torch.tensor([2, 8]))
        self.assertEqual(result_compiled, result_eager)
        self.assertEqual(cnt.frame_count, 1)

        # Pass a contiguous tensor. A recompilation will happen due to 0/1 specialization on stride.
        log_stream, ctx = logs_to_string(
            "torch._functorch._aot_autograd.graph_capture", "aot_graphs"
        )
        with ctx():
            # This used to hit could guard on data-dependent expression Eq(10, u3) x.stride[0]==10. and x.size()=[u2, u3].
            # but not anymore since we use  contiguous_or_false .
            # We need a way to mark strides unbacked to avoid the recompilation here.
            x = torch.randn(10, 10)
            torch._dynamo.decorators.mark_unbacked(x, 0)
            torch._dynamo.decorators.mark_unbacked(x, 1)

        aot_graphs = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            aot_graphs,
            """""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

        result_compiled = compiled_func(x, torch.tensor([2, 50]))
        result_eager = func(x, torch.tensor([2, 50]))

        self.assertEqual(result_compiled, result_eager)
        self.assertEqual(cnt.frame_count, 2)

        x = torch.randn(4, 4)

        result_eager = func(x, torch.tensor([2, 8]))
        result_compiled = compiled_func(x, torch.tensor([2, 8]))
        self.assertEqual(result_compiled, result_eager)
        self.assertEqual(cnt.frame_count, 2)

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_slice(self):
        from torch.fx.experimental.symbolic_shapes import statically_known_true

        # standard slice
        def f1(x, xs):
            u0, u1 = xs.tolist()
            # in this test we add the torch checks not to avoid DDE but to ensure
            # that we pick specific path during compilation.
            torch._check(u0 >= 0)
            torch._check(u0 <= x.size(0))
            torch._check(u1 >= 0)
            torch._check(u1 <= x.size(0))
            torch._check(u0 <= u1)
            out = x[u0:u1]
            if not statically_known_true(out.size(0) == (u1 - u0)):
                raise AssertionError(
                    "expected statically_known_true(out.size(0) == (u1 - u0))"
                )
            return out

        x, xs = torch.randn(10), torch.tensor([3, 6])
        fn1 = torch.compile(f1, fullgraph=True, backend="inductor")
        self.assertEqual(fn1(x, xs).size(0), 3)
        self.assertTrue(torch.allclose(fn1(x, xs), f1(x, xs)))
        with self.assertRaises(RuntimeError):
            fn1(x, torch.tensor([-1, 5]))

        # known negative slice
        def f2(x, n):
            u0 = n.item()
            torch._check(u0 > 1)
            torch._check(u0 <= x.size(0))
            out = x[-u0:]
            if not statically_known_true(out.size(0) == u0):
                raise AssertionError(
                    "expected statically_known_true(out.size(0) == u0)"
                )
            return out

        x, n = torch.randn(10), torch.tensor([5])
        fn2 = torch.compile(f2, fullgraph=True, backend="inductor")
        self.assertEqual(fn2(x, n).size(0), 5)
        self.assertTrue(torch.allclose(fn2(x, n), f2(x, n)))
        with self.assertRaises(RuntimeError):
            fn2(x, torch.tensor([-5]))

        # general case: no known info
        def f3(x, xs):
            u0, u1 = xs.tolist()
            return x[u0:u1]

        log_stream, ctx = logs_to_string(
            "torch._inductor.compile_fx", "post_grad_graphs"
        )
        cnts = CompileCounterWithBackend("inductor")
        x, xs = torch.randn(10), torch.tensor([3, 6])
        with ctx():
            fn3 = torch.compile(f3, fullgraph=True, backend=cnts)
            xs = torch.tensor([-9, -1])  # negative case
            self.assertTrue(torch.allclose(fn3(x, xs), f3(x, xs)))
            xs = torch.tensor([-1000, 1000])  # out of bounds
            self.assertTrue(torch.allclose(fn3(x, xs), f3(x, xs)))
            xs = torch.tensor([2, -2])  # mixed
            self.assertTrue(torch.allclose(fn3(x, xs), f3(x, xs)))
            self.assertEqual(cnts.frame_count, 1)

        aot_graphs = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            aot_graphs,
            """\
        select: "i64[][]cpu" = torch.ops.aten.select.int(arg0_1, 0, 0)
        _local_scalar_dense: "Sym(u0)" = torch.ops.aten._local_scalar_dense.default(select);  select = None
        select_1: "i64[][]cpu" = torch.ops.aten.select.int(arg0_1, 0, 1);  arg0_1 = None
        _local_scalar_dense_1: "Sym(u1)" = torch.ops.aten._local_scalar_dense.default(select_1);  select_1 = None
        slice_1: "f32[u2][1]cpu" = torch.ops.aten.slice.Tensor(arg1_1, 0, _local_scalar_dense, _local_scalar_dense_1);  arg1_1 = _local_scalar_dense = _local_scalar_dense_1 = None
        sym_size_int: "Sym(u2)" = torch.ops.aten.sym_size.int(slice_1, 0)
        ge_1: "Sym(u2 >= 0)" = sym_size_int >= 0
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u2 >= 0 on node 'ge'");  ge_1 = _assert_scalar = None
        le: "Sym(u2 <= 10)" = sym_size_int <= 10;  sym_size_int = None
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(le, "Runtime assertion failed for expression u2 <= 10 on node 'le'");  le = _assert_scalar_1 = None
        sym_storage_offset_default: "Sym(u3)" = torch.ops.aten.sym_storage_offset.default(slice_1)
        ge_2: "Sym(u3 >= 0)" = sym_storage_offset_default >= 0;  sym_storage_offset_default = None
        _assert_scalar_2 = torch.ops.aten._assert_scalar.default(ge_2, "Runtime assertion failed for expression u3 >= 0 on node 'ge_1'");  ge_2 = _assert_scalar_2 = None
        return (slice_1,)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_unbacked_slice_cpp_wrapper(self):
        self.test_unbacked_slice()

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_slice_with_step(self):
        def f1(x, xs):
            u0, u1 = xs.tolist()
            out = x[u0:u1:5]
            return out

        x, xs = torch.randn(10), torch.tensor([2, -2])
        fn1 = torch.compile(f1, fullgraph=True, backend="inductor")
        self.assertTrue(torch.allclose(fn1(x, xs), f1(x, xs)))

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_unbacked_slice_with_step_cpp_wrapper(self):
        self.test_unbacked_slice_with_step()

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_slice_with_tensor_indices(self):
        for d in [True, False]:
            # Test slicing with tensor start/stop/step on RHS (reading)

            # Test 1: Basic slice with tensor start and stop
            def f1(x, start_t, stop_t):
                return x[start_t:stop_t]

            x = torch.randn(20)
            start_t = torch.tensor(5)
            stop_t = torch.tensor(15)
            fn1 = torch.compile(f1, fullgraph=True, dynamic=d, backend="inductor")
            self.assertTrue(
                torch.allclose(fn1(x, start_t, stop_t), f1(x, start_t, stop_t))
            )

            # Test 2: Slice with tensor step
            def f2(x, start_t, stop_t, step_t):
                return x[start_t:stop_t:step_t]

            step_t = torch.tensor(2)
            fn2 = torch.compile(f2, fullgraph=True, dynamic=d, backend="inductor")
            self.assertTrue(
                torch.allclose(
                    fn2(x, start_t, stop_t, step_t), f2(x, start_t, stop_t, step_t)
                )
            )

            # Test 3: Slice with only tensor start
            def f3(x, start_t):
                return x[start_t:]

            fn3 = torch.compile(f3, fullgraph=True, dynamic=d, backend="inductor")
            self.assertTrue(torch.allclose(fn3(x, start_t), f3(x, start_t)))

            # Test 4: Slice with only tensor stop
            def f4(x, stop_t):
                return x[:stop_t]

            fn4 = torch.compile(f4, fullgraph=True, dynamic=d, backend="inductor")
            self.assertTrue(torch.allclose(fn4(x, stop_t), f4(x, stop_t)))

            # Test 5: Negative indices with tensors
            def f5(x, start_t):
                return x[start_t:-1]

            start_t_neg = torch.tensor(-10)
            fn5 = torch.compile(f5, fullgraph=True, dynamic=d, backend="inductor")
            self.assertTrue(torch.allclose(fn5(x, start_t_neg), f5(x, start_t_neg)))

            # Test 6: Multidimensional slice with tensor indices
            def f6(x, start_t, stop_t):
                return x[:, start_t:stop_t]

            x_2d = torch.randn(10, 20)
            fn6 = torch.compile(f6, fullgraph=True, dynamic=d, backend="inductor")
            self.assertTrue(
                torch.allclose(fn6(x_2d, start_t, stop_t), f6(x_2d, start_t, stop_t))
            )

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_slice_with_tensor_indices_cpp_wrapper(self):
        self.test_slice_with_tensor_indices()

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_tensor_split(self):
        def f1(x, xs):
            xs = torch.tensor(xs.tolist())
            return torch.tensor_split(x, xs)

        x = torch.randn(20)
        xs = torch.tensor([5, 10, 15])
        fn = torch.compile(f1, fullgraph=True, backend="inductor")

        def compare(x, xs):
            for i, j in zip(f1(x, xs), fn(x, xs)):
                self.assertTrue(torch.allclose(i, j))

        compare(x, xs)
        xs = torch.tensor([-15, 9, 10, 11])
        compare(x, xs)
        xs = torch.tensor([-15, -10, -5, -2])
        compare(x, xs)

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_tensor_split_cpp_wrapper(self):
        self.test_tensor_split()

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._dynamo.config.patch("capture_dynamic_output_shape_ops", True)
    def test_nonzero_slice(self):
        def f(x):
            nz = x.nonzero()
            return nz[:-1]

        x = torch.randn(3, 4)
        fn = torch.compile(f, fullgraph=True, backend="inductor")
        self.assertTrue(torch.allclose(f(x), fn(x)))
        y = torch.zeros(3, 4)
        self.assertTrue(torch.allclose(f(y), fn(y)))

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._dynamo.config.patch("capture_dynamic_output_shape_ops", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_nonzero_slice_cpp_wrapper(self):
        self.test_nonzero_slice()

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._dynamo.config.patch("capture_dynamic_output_shape_ops", True)
    def test_nonzero_select(self):
        def f(x):
            nz = x.nonzero()
            return nz[-1] + nz[0]

        x = torch.randn(3, 4)
        fn = torch.compile(f, fullgraph=True, backend="inductor")
        self.assertTrue(torch.allclose(f(x), fn(x)))

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._dynamo.config.patch("capture_dynamic_output_shape_ops", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_nonzero_select_cpp_wrapper(self):
        self.test_nonzero_select()

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._dynamo.config.patch("capture_dynamic_output_shape_ops", True)
    def test_padnd(self):
        import torch.nn.functional as F

        def f(x, xs, y):
            u0, u1 = xs.tolist()
            for u in [u0, u1]:
                torch._check(u >= 0)
            z = F.pad(x, (u0, u1, u0, u1))
            return z @ y

        x = torch.randn(8, 8)
        xs = torch.tensor([2, 2])
        y = torch.randn(12, 4)
        fn = torch.compile(f, fullgraph=True, backend="inductor")
        fn(x, xs, y)

    @unittest.skip("this test fails due to inductor/autograd issue #153041")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_non_contigious_reshape_failing(self):
        # reshape u1 -> (u0*u0)
        # this result in the tensor "i64[u0, u0][s7*u0, s7].
        # reshape happens in place reshape (no-clone)
        def func(x, y):
            f = y.item()
            t1 = x.view((f, f))
            t2 = x.reshape((f, f))
            return t1, t2

        # create a non-contiguous with data being even numbers in [0:cnt-1]
        def make_non_contiguous_tensor(cnt):
            # create a non-contiguous tensor x that is skipping odd indices.
            x = torch.arange(cnt * 2)
            x = x.as_strided((x.size()[0] // 2,), (2,))
            return x

        x = make_non_contiguous_tensor(4)
        torch._dynamo.decorators.mark_unbacked(x, 0)
        compiled_func = torch.compile(
            fullgraph=True,
            backend="inductor",
        )(func)

        compiled_result = compiled_func(x, torch.tensor([2]))
        eager_result = func(x, torch.tensor([2]))
        self.assertEqual(compiled_result, eager_result)

    @skipIfTorchDynamo("not allowed to trace mark_unbacked")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_invalid_view_unbacked_view(self):
        cnt = CompileCounterWithBackend("inductor")

        # This view (u2, u3) -> (u0, u1) can't happen in general unless we know that input is contiguous or we have
        # hints to to compute strides.
        def func(x, y):
            u0, u1 = y.tolist()
            result2 = x.view(u0, u1) * 10
            return result2

        compiled_func = torch.compile(fullgraph=True, backend=cnt, dynamic=True)(func)

        x = torch.randn(10, 10)
        # make x not contiguous.
        x = x.t_()
        torch._dynamo.decorators.mark_unbacked(x, 0)
        torch._dynamo.decorators.mark_unbacked(x, 1)
        with self.assertRaises(torch._dynamo.exc.UserError):
            # throws a data dependent error.
            compiled_func(x, torch.tensor([5, 20]))

    @skipIfTorchDynamo()
    def test_unbind_not_dynamic(self):
        cnt = CompileCounter()

        @torch.compile(fullgraph=True, dynamic=True, backend=cnt)
        def func(y):
            return y.unbind(dim=2), y * 10

        func(torch.ones(5, 6, 7, 8))
        self.assertEqual(cnt.frame_count, 1)
        # it can be dynamic in all dimensions except dim=2
        func(torch.ones(4, 9, 7, 10))
        self.assertEqual(cnt.frame_count, 1)

        func(torch.ones(5, 6, 8, 8))
        func(torch.ones(5, 6, 9, 8))
        self.assertEqual(cnt.frame_count, 3)

    @skipIfTorchDynamo("not allowed to trace mark_unbacked")
    @fresh_cache()
    def test_unbacked_contiguous(self):
        cnt = CompileCounterWithBackend("inductor")

        def func(x):
            contig = x.contiguous()
            return (contig + 1) * 100

        compiled_func = torch.compile(fullgraph=True, backend=cnt, dynamic=True)(func)

        x = torch.randn(10, 10)
        # make x not contiguous.
        x = x.t_()
        torch._dynamo.decorators.mark_unbacked(x, 0)
        torch._dynamo.decorators.mark_unbacked(x, 1)
        log_stream, ctx = logs_to_string(
            "torch._inductor.compile_fx", "post_grad_graphs"
        )
        with ctx():
            compiled_func(x)
            self.assertEqual(compiled_func(x), func(x))
            y = torch.rand(20, 20).t()
            self.assertEqual(compiled_func(y), func(y))
            self.assertEqual(cnt.frame_count, 1)
        output = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            output,
            """\
        ge: "Sym(u0 >= 0)" = arg0_1 >= 0;  arg0_1 = None
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u0 >= 0 on node 'ge'");  ge = _assert_scalar = None
        ge_1: "Sym(u1 >= 0)" = arg1_1 >= 0;  arg1_1 = None
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u1 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
        clone: "f32[u0, u1][Max(1, u1), 1]cpu" = torch.ops.aten.clone.default(arg2_1, memory_format = torch.contiguous_format);  arg2_1 = None
        add_3: "f32[u0, u1][Max(1, u1), 1]cpu" = torch.ops.aten.add.Tensor(clone, 1);  clone = None
        mul_6: "f32[u0, u1][Max(1, u1), 1]cpu" = torch.ops.aten.mul.Tensor(add_3, 100);  add_3 = None
        return (mul_6,)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

        log_stream, ctx = logs_to_string(
            "torch._inductor.compile_fx", "post_grad_graphs"
        )
        with ctx():
            # recompilation will happen due to stride specialization.
            y = torch.rand(20, 20)
            torch._dynamo.decorators.mark_unbacked(y, 0)
            torch._dynamo.decorators.mark_unbacked(y, 1)
            self.assertEqual(compiled_func(y), func(y))
            self.assertEqual(cnt.frame_count, 2)

        output = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()

        # No clone this time since input is contiguous.
        self.assertExpectedInline(
            output,
            """\
        ge: "Sym(u0 >= 0)" = arg0_1 >= 0;  arg0_1 = None
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u0 >= 0 on node 'ge'");  ge = _assert_scalar = None
        ge_1: "Sym(u1 >= 0)" = arg1_1 >= 0;  arg1_1 = None
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u1 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
        add: "f32[u0, u1][Max(1, u1), 1]cpu" = torch.ops.aten.add.Tensor(arg2_1, 1);  arg2_1 = None
        mul_5: "f32[u0, u1][Max(1, u1), 1]cpu" = torch.ops.aten.mul.Tensor(add, 100);  add = None
        return (mul_5,)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_select_index(self):
        cnt = CompileCounterWithBackend("inductor")

        def func(x, y):
            u0 = y.item()
            return (
                torch.select(x, 0, u0),
                torch.select(x, 1, u0),
                torch.select(x, 2, u0),
            )

        compiled_func = torch.compile(fullgraph=True, backend=cnt, dynamic=True)(func)
        x = torch.rand(3, 3, 3)
        zero = torch.tensor([0])
        pos = torch.tensor([1])
        # code can handle both negative and positive indices.
        neg = torch.tensor([-1])

        log_stream, ctx = logs_to_string(
            "torch._inductor.compile_fx", "post_grad_graphs"
        )
        with ctx():
            self.assertEqual(compiled_func(x, zero), func(x, zero))
        output = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            output,
            """\
        _local_scalar_dense: "Sym(u0)" = torch.ops.aten._local_scalar_dense.default(arg0_1);  arg0_1 = None
        select: "f32[s77, s77][s77, 1]cpu" = torch.ops.aten.select.int(arg2_1, 0, _local_scalar_dense)
        select_1: "f32[s77, s77][s77**2, 1]cpu" = torch.ops.aten.select.int(arg2_1, 1, _local_scalar_dense)
        select_2: "f32[s77, s77][s77**2, s77]cpu" = torch.ops.aten.select.int(arg2_1, 2, _local_scalar_dense);  arg2_1 = _local_scalar_dense = None
        return (select, select_1, select_2)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )
        self.assertEqual(compiled_func(x, pos), func(x, pos))
        self.assertEqual(compiled_func(x, neg), func(x, neg))
        self.assertEqual(cnt.frame_count, 1)

        def func2(x, y):
            u0, u1 = y.tolist()
            return torch.select(x, 0, u0 + u1)

        compiled_func2 = torch.compile(fullgraph=True, backend=cnt, dynamic=False)(
            func2
        )
        zero = torch.tensor([0, 0])
        pos = torch.tensor([1, 1])
        neg = torch.tensor([-1, -1])

        self.assertEqual(compiled_func2(x, pos), func2(x, pos))
        self.assertEqual(compiled_func2(x, neg), func2(x, neg))
        self.assertEqual(compiled_func2(x, zero), func2(x, zero))
        self.assertEqual(cnt.frame_count, 2)

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_select_2(self):
        class M(torch.nn.Module):
            def forward(self, x):
                nz = x.nonzero()
                return nz[-1]

        mod = M()
        x = torch.randn(4)
        self.assertEqual(torch.compile(mod)(x), mod(x))

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_select_index_with_check(self):
        def func3(x, y):
            u0 = y.item()
            # Test that taking the non-unbacked path works fine also.
            torch._check(u0 >= 0)
            return (torch.select(x, 1, u0),)

        compiled_func3 = torch.compile(
            fullgraph=True, backend="inductor", dynamic=True
        )(func3)
        x = torch.rand(3, 3, 3)
        zero = torch.tensor([0])
        pos = torch.tensor([1])
        print(compiled_func3(x, pos))

        self.assertEqual(compiled_func3(x, pos), func3(x, pos))
        self.assertEqual(compiled_func3(x, zero), func3(x, zero))

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_unbacked_select_index_cpp_wrapper(self):
        self.test_unbacked_select_index()

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_select2(self):
        def f(idx, x):
            x = x.select(0, idx.item())
            return x @ x

        x = torch.randn(3, 3, 3)
        idx = torch.tensor(1, dtype=torch.int64)
        out = torch.compile(f)(idx, x)
        self.assertEqual(out, f(idx, x))

    def test_trunc_int_div_true(self):
        @torch.compile(backend="inductor", dynamic=True, fullgraph=True)
        def f(x, s13, s57, s77):
            torch._check(s13 >= 0)
            torch._check(s57 >= 0)
            torch._check(s77 >= 0)
            if int(s13 * ((s57 // s13) + (s77 // s13)) / s13) >= 1:
                return x * 2
            else:
                return x * 100

        # ensure we compile this with no errors.
        x = torch.rand(10)
        f(x, 4, 4096, 3920)

    @skipIfTorchDynamo("not allowed to trace mark_unbacked")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_reshape3(self):
        def func(x):
            x = x.as_strided([x.size()[0], 1536], [2048, 1])
            result1 = x.view(x.size()[0], -1, 128)
            return result1 * 10

        compiled = torch.compile(fullgraph=True, backend="inductor")(func)
        x = torch.randn(10, 2048)

        torch._dynamo.decorators.mark_unbacked(x, 0)
        self.assertEqual(func(x), compiled(x))

    @fresh_cache()
    @skipIfTorchDynamo("not allowed to trace mark_unbacked")
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_reshape_copy(self):
        cnt = CompileCounterWithBackend("inductor")

        # Reshape happens in place reshape (no-clone)
        # reshape u1 -> (u0*u0)
        def func(x, y):
            f = y.item()
            t3 = torch._ops.ops.aten._reshape_copy(x, (f, f))
            return t3

        compiled_func = torch.compile(
            fullgraph=True,
            backend=cnt,
            dynamic=True,
        )(func)

        # create a non-contiguous with data being even numbers in [0:cnt-1]
        # and reshape it into sqrt(cnt)*sqrt(cnt)
        def make_non_contiguous_tensor_and_test(cnt):
            # create a non-contiguous tensor x that is skipping odd indices.
            x = torch.arange(cnt * 2)
            x = x.as_strided((x.size()[0] // 2,), (2,))

            torch._dynamo.decorators.mark_unbacked(x, 0)
            sz = torch.tensor([int(math.sqrt(cnt))])
            compiled_result = compiled_func(x, sz)
            eager_result = func(x, sz)
            self.assertEqual(compiled_result, eager_result)

        log_stream, ctx = logs_to_string(
            "torch._functorch._aot_autograd.graph_capture", "aot_graphs"
        )
        with ctx():
            make_non_contiguous_tensor_and_test(4)
        aot_graphs = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            aot_graphs,
            """\
def forward(self, arg0_1: "i64[1][1]cpu", arg1_1: "Sym(u1)", arg2_1: "Sym(s7)", arg3_1: "i64[u1][s7]cpu"):
        ge: "Sym(u1 >= 0)" = arg1_1 >= 0
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u1 >= 0 on node 'ge'");  ge = _assert_scalar = None
        _local_scalar_dense: "Sym(u0)" = torch.ops.aten._local_scalar_dense.default(arg0_1);  arg0_1 = None
        ge_1: "Sym(u0 >= 0)" = _local_scalar_dense >= 0
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u0 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
        pow_1: "Sym(u0**2)" = _local_scalar_dense ** 2
        eq: "Sym(Eq(u1, u0**2))" = arg1_1 == pow_1;  arg1_1 = pow_1 = None
        _assert_scalar_2 = torch.ops.aten._assert_scalar.default(eq, "Runtime assertion failed for expression Eq(u1, u0**2) on node 'eq'");  eq = _assert_scalar_2 = None
        _reshape_copy: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten._reshape_copy.default(arg3_1, [_local_scalar_dense, _local_scalar_dense]);  arg3_1 = _local_scalar_dense = None
        return (_reshape_copy,)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

        make_non_contiguous_tensor_and_test(49)
        self.assertEqual(cnt.frame_count, 1)

        # Pass in a contiguous tensor, it will recompile due to stride being 1 (0/1 specialization).
        # marking strides unbacked would have avoided the recompilation here.
        x = torch.arange(100)
        torch._dynamo.decorators.mark_unbacked(x, 0)

        log_stream, ctx = logs_to_string(
            "torch._functorch._aot_autograd.graph_capture", "aot_graphs"
        )
        with ctx():
            compiled_result = compiled_func(x, torch.tensor([10]))
            eager_result = func(x, torch.tensor([10]))
            self.assertEqual(compiled_result, eager_result)
            self.assertEqual(cnt.frame_count, 2)

        aot_graphs = "\n".join(log_stream.getvalue().strip().split("\n")[4:]).strip()
        self.assertExpectedInline(
            aot_graphs,
            """\
def forward(self, arg0_1: "i64[1][1]cpu", arg1_1: "Sym(u1)", arg2_1: "i64[u1][1]cpu"):
        ge: "Sym(u1 >= 0)" = arg1_1 >= 0
        _assert_scalar = torch.ops.aten._assert_scalar.default(ge, "Runtime assertion failed for expression u1 >= 0 on node 'ge'");  ge = _assert_scalar = None
        _local_scalar_dense: "Sym(u0)" = torch.ops.aten._local_scalar_dense.default(arg0_1);  arg0_1 = None
        ge_1: "Sym(u0 >= 0)" = _local_scalar_dense >= 0
        _assert_scalar_1 = torch.ops.aten._assert_scalar.default(ge_1, "Runtime assertion failed for expression u0 >= 0 on node 'ge_1'");  ge_1 = _assert_scalar_1 = None
        pow_1: "Sym(u0**2)" = _local_scalar_dense ** 2
        eq: "Sym(Eq(u1, u0**2))" = arg1_1 == pow_1;  arg1_1 = pow_1 = None
        _assert_scalar_2 = torch.ops.aten._assert_scalar.default(eq, "Runtime assertion failed for expression Eq(u1, u0**2) on node 'eq'");  eq = _assert_scalar_2 = None
        _reshape_copy: "i64[u0, u0][Max(1, u0), 1]cpu" = torch.ops.aten._reshape_copy.default(arg2_1, [_local_scalar_dense, _local_scalar_dense]);  arg2_1 = _local_scalar_dense = None
        return (_reshape_copy,)""",  # noqa: B950
            ignore_comments=True,
            ignore_empty_lines=True,
        )

        x = torch.arange(25)
        compiled_result = compiled_func(x, torch.tensor([5]))
        eager_result = func(x, torch.tensor([5]))
        self.assertEqual(cnt.frame_count, 2)

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_item(self):
        def func():
            _x_ms = torch.tensor([True, False], dtype=torch.int64)
            _mask_ms = torch.zeros_like(_x_ms, dtype=torch.bool)
            _mask_ms[:1] = True
            var_node_2 = torch.masked_select(_x_ms, _mask_ms)
            var_node_0 = var_node_2.item()
            return var_node_0

        result_original = func()
        compiled_program = torch.compile(func, fullgraph=True, dynamic=True)
        result_compiled = compiled_program()
        self.assertEqual(result_original, result_compiled)

    def test_unbacked_item_set_item(self):
        def my_arithmetic(a, b):
            wrk = torch.zeros(a.size(0))
            for i in range(a.size(0)):
                idx = b[i].item()
                wrk[idx] += 1

            return wrk

        compiled = torch.compile(my_arithmetic, fullgraph=True, disable=False)
        a = torch.randn([9])
        b = torch.ones(9, dtype=torch.int32)
        compiled(a, b)
        self.assertEqual(compiled(a, b), my_arithmetic(a, b))

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_item_set_item2(self):
        def accumulate(X0, start):
            start = start.item()
            N = 3
            result = X0[start]
            for i in range(N):
                result += X0[start + 1 + i]
            return result

        compiled = torch.compile(accumulate, fullgraph=True)
        X0 = torch.randn(10, 10)
        self.assertEqual(
            accumulate(X0, torch.tensor([1])), compiled(X0, torch.tensor([1]))
        )

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_unbacked_item_set_item3(self):
        def func(x, y):
            u0 = y.item()
            x[u0] = 0
            return x

        compiled = torch.compile(func, fullgraph=True, disable=False)
        b = torch.tensor([0])
        a = torch.ones(9, dtype=torch.int32)

        compiled(a, b)
        self.assertEqual(compiled(a, b), func(a, b))

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_select_scatter_unbacked_index(self):
        def func(x, y):
            u0 = y.item()
            # Create a scalar tensor to scatter into the selected index
            scalar_src = torch.tensor(42, dtype=x.dtype)
            return x.select_scatter(scalar_src, 0, u0)

        compiled = torch.compile(func, fullgraph=True, dynamic=True, backend="inductor")
        b = torch.tensor([0])
        a = torch.ones(9, dtype=torch.int32)

        self.assertEqual(compiled(a, b), func(a, b))

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_narrow_unbacked_start(self):
        def func(x, start, length):
            # unbacked start
            u0 = start.item()
            return torch.narrow(x, 0, u0, length)

        compiled_func = torch.compile(func, fullgraph=True, backend="inductor")

        x = torch.tensor([1, 2, 3, 4, 5, 6])

        # Test cases: (start, length)
        test_cases = [
            # Negative starts
            (-2, 2),  # Start from second-to-last element
            (-1, 1),  # Start from last element
            (-3, 3),  # Start from third-to-last element
            (-6, 2),  # Start from beginning (negative)
            (-4, 1),  # Start from fourth-to-last element
            # Positive starts
            (0, 2),  # Start from beginning
            (1, 3),  # Start from second element
            (2, 2),  # Start from third element
            (4, 2),  # Start near end
            # Edge cases
            (0, 6),  # Full tensor
            (0, 1),  # Single element from start
            (5, 1),  # Single element from end
        ]

        for start_val, length in test_cases:
            with self.subTest(start=start_val, length=length):
                start = torch.tensor([start_val])

                # Test with compiled function
                result_compiled = compiled_func(x, start, length)

                # Test with eager function (expected behavior)
                result_eager = func(x, start, length)

                # Compare results
                self.assertEqual(result_compiled, result_eager)

    @fresh_cache()
    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    @torch._inductor.config.patch("cpp_wrapper", True)
    def test_narrow_unbacked_start_cpp_wrapper(self):
        """Test narrow with unbacked start with cpp_wrapper"""
        self.test_narrow_unbacked_start()

    @torch._dynamo.config.patch(capture_scalar_outputs=True)
    def test_narrow_with_tensor_start(self):
        @torch.compile(backend="inductor", fullgraph=True)
        def f(x, start, end):
            return torch.narrow(x, 0, start, end)

        x = torch.tensor(
            [False], device="cuda:0" if torch.cuda.is_available() else "cpu"
        )
        start = torch.tensor(0)
        res = f(x, start, 0)
        self.assertEqual(res.shape, torch.Size([0]))

    @skipIfTorchDynamo("mark_unbacked not supported")
    def test_unbacked_norm_no_dde(self):
        def vector_norm(x):
            return torch.linalg.vector_norm(x)

        def dist_fn(x, y):
            return torch.dist(x, y)

        def norm_fro(x):
            return torch.norm(x, p="fro")

        x = torch.randn(4, 5)
        y = torch.randn(4, 5)
        for dim in range(x.ndim):
            torch._dynamo.decorators.mark_unbacked(x, dim)
            torch._dynamo.decorators.mark_unbacked(y, dim)

        for fn, args in [
            (vector_norm, (x,)),
            (dist_fn, (x, y)),
            (norm_fro, (x,)),
        ]:
            torch._dynamo.reset()
            compiled = torch.compile(fn, fullgraph=True, backend="eager")
            compiled(*args)

    @skipIfTorchDynamo()
    @torch.fx.experimental._config.patch("backed_size_oblivious", True)
    def test_backed_size_oblivious_expand(self):
        cnt = CompileCounterWithBackend("inductor")
        torch._dynamo.reset()

        def func(a, shape):
            return a.expand(*shape)

        compiled = torch.compile(func, fullgraph=True, backend=cnt, dynamic=True)

        def run(a, shape):
            self.assertEqual(compiled(a, shape), func(a, shape))

        # No specialization - matching dimensions
        run(torch.rand(2, 10), (2, 10))
        run(torch.rand(5, 5), (5, 5))
        self.assertEqual(cnt.frame_count, 1)  # Single compilation

        cnt.clear()
        torch._dynamo.reset()

        # Specialize input dim to 1 (broadcasting)
        # Input dim is 1, needs to expand to larger size
        run(torch.rand(1, 10), (9, 10))  # Compile with input[0] == 1
        run(torch.rand(1, 5), (7, 5))  # Reuse same compiled graph
        self.assertEqual(cnt.frame_count, 1)
        # Changing input from 1 triggers recompilation
        run(torch.rand(5, 10), (5, 10))
        self.assertEqual(cnt.frame_count, 2)

        cnt.clear()
        torch._dynamo.reset()

        # Multi-dimensional broadcasting
        # Multiple dimensions with different broadcast semantics
        run(torch.rand(1, 1, 10), (5, 7, 10))  # Broadcast first two dims
        run(torch.rand(1, 1, 5), (3, 4, 5))  # Reuse pattern
        self.assertEqual(cnt.frame_count, 1)

    @skipIfTorchDynamo()
    @torch.fx.experimental._config.patch("backed_size_oblivious", True)
    def test_backed_size_oblivious_broadcast(self):
        cnt = CompileCounterWithBackend("inductor")
        torch._dynamo.reset()

        def func(a, b):
            torch.broadcast_shapes(a.size(), b.size())
            return a + b

        compiled = torch.compile(func, fullgraph=True, backend=cnt, dynamic=True)

        def run(a, b):
            self.assertEqual(compiled(a, b), func(a, b))

        # No 0/1 specializations, no broadcasts.
        # but a[0] == b[0] and a[1] == b[1] are asserted.
        run(torch.rand(1, 10), torch.rand(1, 10))
        run(torch.rand(1, 1), torch.rand(1, 1))
        run(torch.rand(10, 10), torch.rand(10, 10))

        self.assertEqual(cnt.frame_count, 1)
        run(torch.rand(10, 10), torch.rand(1, 10))
        self.assertEqual(cnt.frame_count, 2)

        cnt.clear()
        torch._dynamo.reset()

        # specialize a[0] == 1. b[0] not specialized.
        run(torch.rand(1, 10), torch.rand(9, 10))
        run(torch.rand(1, 10), torch.rand(1, 10))
        self.assertEqual(cnt.frame_count, 1)
        # if we change a[0] we get recompilation.
        run(torch.rand(10, 10), torch.rand(10, 10))
        self.assertEqual(cnt.frame_count, 2)

        cnt.clear()
        torch._dynamo.reset()

        # TODO duck sizing shall be disabled when backed_size_oblivious
        # is on probably.
        # specialize b[0] == 1. a[0] not specialized.
        run(torch.rand(10, 11), torch.rand(1, 11))
        run(torch.rand(1, 10), torch.rand(1, 10))
        self.assertEqual(cnt.frame_count, 1)
        run(torch.rand(2, 10), torch.rand(2, 10))
        self.assertEqual(cnt.frame_count, 2)

    @torch._dynamo.config.patch("capture_dynamic_output_shape_ops", True)
    def test_unbacked_view_extra(self):
        def fn(x):
            i0 = x.nonzero().size(0)
            y = torch.zeros((i0, 192))
            return y.view([12, -1, 192])

        res1 = torch.compile(fn, fullgraph=True)(torch.ones((12,)))
        res2 = fn(torch.ones((12,)))
        self.assertEqual(res1, res2)

    @skipIfTorchDynamo("mark_dynamic not supported")
    def test_hint_override_consistent_stride1(self):
        @torch.compile(fullgraph=True, dynamic=True)
        def func(x):
            a = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                x.size()[2]
            )
            b = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                x.stride()[1]
            )
            torch._check(a == b)
            torch._check(a == 6)

            a = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                x.size()[1] * x.size()[2]
            )
            b = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                x.stride()[0]
            )
            torch._check(a == b)
            torch._check(a == 120)

            return a, b, x * 100

        x = torch.rand(10, 20, 30)

        torch._dynamo.mark_dynamic(x, 0, hint_override=2)
        torch._dynamo.mark_dynamic(x, 2, hint_override=6)

        func(x)

    @skipIfTorchDynamo("mark_dynamic not supported")
    def test_hint_override_consistent_stride2(self):
        @torch.compile(fullgraph=True, dynamic=True)
        def func(x):
            # only one of the sizes has hint overridden.
            a = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                x.size()[1] * x.size()[2]
            )
            b = torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                x.stride()[0]
            )
            torch._check(a == b)
            torch._check(a == 24)

            return a, b, x * 100

        x = torch.rand(10, 20, 30)

        torch._dynamo.mark_dynamic(x, 0, hint_override=2)
        torch._dynamo.mark_dynamic(x, 1, hint_override=4)
        torch._dynamo.mark_dynamic(x, 2, hint_override=6)
        func(x)

    def test_optimization_hint(self):
        @torch.compile(fullgraph=True)
        def func(x):
            u0 = x.item()
            a = torch.ones([u0])
            torch._check(
                torch.fx.experimental.symbolic_shapes.optimization_hint(
                    a.size()[0], fallback=300
                )
                == 300
            )
            b = torch.ones([x.item() * 2])
            torch._check(
                torch.fx.experimental.symbolic_shapes.optimization_hint(
                    b.size()[0], fallback=300
                )
                == 600
            )

            return a * 10

        func(torch.tensor([33]))

    def test_guarding_hint_or_throw(self):
        @torch.compile(fullgraph=True)
        def func(x):
            u0 = x.item()
            a = torch.ones([u0])
            torch._check(
                torch.fx.experimental.symbolic_shapes.guarding_hint_or_throw(
                    a.size()[0]
                )
                == 300
            )

            return a * 10

        self.assertRaises(GuardOnDataDependentSymNode, lambda: func(torch.tensor([33])))

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_shape_id_unifies_unbacked_symbols(self):
        """
        Test that shape_id parameter in mark_unbacked causes tensors to share
        the same unbacked symbol, allowing equality comparisons without DDE.
        """

        def func(x, y):
            if x.size()[0] == y.size()[0]:
                return x + y
            else:
                return x - y

        # Test 1: Without shape_id, comparing unbacked symbols raises DDE
        x1 = torch.rand(4, 3)
        y1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0)
        torch._dynamo.decorators.mark_unbacked(y1, 0)

        torch._dynamo.reset()
        with self.assertRaises(torch._dynamo.exc.UserError):
            compiled_func = torch.compile(func, fullgraph=True, backend="eager")
            compiled_func(x1, y1)

        # Test 2: With shape_id, same symbol is used - no DDE, fullgraph succeeds
        x2 = torch.rand(4, 3)
        y2 = torch.rand(4, 3)

        torch._dynamo.reset()
        compiled_func = torch.compile(func, fullgraph=True, backend="eager")
        result = compiled_func(x2, y2)
        # Since both use same symbol, x.size()[0] == y.size()[0] is always True
        self.assertTrue(torch.allclose(result, x2 + y2))

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_shape_id_runtime_assertion_on_mismatch(self):
        """
        Test that shape_id with different actual sizes at runtime
        raises an assertion error during tracing. This ensures that duck shaping
        violations are caught rather than silently producing incorrect results.
        """

        def func(x, y):
            return x + y

        torch._dynamo.reset()
        compiled_func = torch.compile(func, fullgraph=True, backend="eager")

        # First, compile with valid inputs (same batch size)
        x1 = torch.rand(4, 3)
        y1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0, shape_id="batch")
        torch._dynamo.decorators.mark_unbacked(y1, 0, shape_id="batch")
        result = compiled_func(x1, y1)
        self.assertTrue(torch.allclose(result, x1 + y1))

        # Now pass tensors with different batch sizes but same shape_id
        # This triggers recompilation, and during tracing the torch._check
        # equality assertion will fail because the sizes don't match
        x2 = torch.rand(4, 3)
        y2 = torch.rand(5, 3)  # Different batch size!
        torch._dynamo.decorators.mark_unbacked(x2, 0, shape_id="batch")
        torch._dynamo.decorators.mark_unbacked(y2, 0, shape_id="batch")

        # Should raise an AssertionError during guard building because batch sizes don't match
        with self.assertRaises(AssertionError):
            compiled_func(x2, y2)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_shape_id_recompilation(self):
        """
        Test that changing _dynamo_shape_ids triggers recompilation.
        """
        counter = CompileCounter()

        def func(x):
            return x + 1

        compiled_func = torch.compile(func, backend=counter)

        # First call with shape_id
        x1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0, shape_id="batch")
        compiled_func(x1)
        self.assertEqual(counter.frame_count, 1)

        # Second call with same shape_id - no recompilation
        x2 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x2, 0, shape_id="batch")
        compiled_func(x2)
        self.assertEqual(counter.frame_count, 1)

        # Third call without shape_id - should recompile
        x3 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x3, 0)
        compiled_func(x3)
        self.assertEqual(counter.frame_count, 2)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_shape_id_no_recompile_without_unbacked_indices(self):
        """
        Test that passing a tensor without _dynamo_unbacked_indices after
        compiling with shape_ids does NOT trigger recompilation.
        The guard on shape_ids only applies when the runtime tensor
        also has _dynamo_unbacked_indices.
        """
        counter = CompileCounter()

        def func(x):
            return x + 1

        compiled_func = torch.compile(func, backend=counter)

        # First call with shape_id (has _dynamo_unbacked_indices)
        x1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0, shape_id="batch")
        compiled_func(x1)
        self.assertEqual(counter.frame_count, 1)

        # Second call with regular tensor (no _dynamo_unbacked_indices)
        # Should NOT recompile - guard passes when no _dynamo_unbacked_indices
        x2 = torch.rand(4, 3)
        compiled_func(x2)
        self.assertEqual(counter.frame_count, 1)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_shape_id_recompile_with_different_id(self):
        """
        Test that passing a tensor with same _dynamo_dynamic_indices but
        different shape_id DOES trigger recompilation.
        """
        counter = CompileCounter()

        def func(x):
            return x + 1

        compiled_func = torch.compile(func, backend=counter)

        # First call with shape_id="batch"
        x1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0, shape_id="batch")
        compiled_func(x1)
        self.assertEqual(counter.frame_count, 1)

        # Second call with different shape_id - should recompile
        x2 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x2, 0, shape_id="other")
        compiled_func(x2)
        self.assertEqual(counter.frame_count, 2)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_unbacked_bounds_recompilation(self):
        """
        Test that changing _dynamo_unbacked_bounds triggers recompilation.
        """
        counter = CompileCounter()

        def func(x):
            return x + 1

        compiled_func = torch.compile(func, backend=counter)

        # First call with min/max bounds
        x1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0, min=1, max=100)
        compiled_func(x1)
        self.assertEqual(counter.frame_count, 1)

        # Second call with same bounds - no recompilation
        x2 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x2, 0, min=1, max=100)
        compiled_func(x2)
        self.assertEqual(counter.frame_count, 1)

        # Third call without bounds - should recompile
        x3 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x3, 0)
        compiled_func(x3)
        self.assertEqual(counter.frame_count, 2)

        # Fourth call with different bounds - should recompile
        x4 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x4, 0, min=1, max=200)
        compiled_func(x4)
        self.assertEqual(counter.frame_count, 3)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_unbacked_bounds_no_recompile_without_unbacked_indices(self):
        """
        Test that passing a tensor without _dynamo_unbacked_indices after
        compiling with bounds does NOT trigger recompilation.
        The guard on bounds only applies when the runtime tensor
        also has _dynamo_unbacked_indices.
        """
        counter = CompileCounter()

        def func(x):
            return x + 1

        compiled_func = torch.compile(func, backend=counter)

        # First call with bounds (has _dynamo_unbacked_indices)
        x1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0, min=1, max=100)
        compiled_func(x1)
        self.assertEqual(counter.frame_count, 1)

        # Second call with regular tensor (no _dynamo_unbacked_indices)
        # Should NOT recompile - guard passes when no _dynamo_unbacked_indices
        x2 = torch.rand(4, 3)
        compiled_func(x2)
        self.assertEqual(counter.frame_count, 1)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_unbacked_no_bounds_then_bounds(self):
        """
        Test that compiling without bounds then calling with bounds
        triggers recompilation.
        """
        counter = CompileCounter()

        def func(x):
            return x + 1

        compiled_func = torch.compile(func, backend=counter)

        # First call with mark_unbacked but no bounds
        x1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0)
        compiled_func(x1)
        self.assertEqual(counter.frame_count, 1)

        # Second call with same (no bounds) - no recompilation
        x2 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x2, 0)
        compiled_func(x2)
        self.assertEqual(counter.frame_count, 1)

        # Third call with min=1, max=100 - should recompile
        # (compiled without bounds, now calling with bounds)
        x3 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x3, 0, min=1, max=100)
        compiled_func(x3)
        self.assertEqual(counter.frame_count, 2)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_unbacked_no_shape_id_then_shape_id(self):
        """
        Test that compiling without shape_id then calling with shape_id
        triggers recompilation.
        """
        counter = CompileCounter()

        def func(x):
            return x + 1

        compiled_func = torch.compile(func, backend=counter)

        # First call with mark_unbacked but no shape_id
        x1 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x1, 0)
        compiled_func(x1)
        self.assertEqual(counter.frame_count, 1)

        # Second call with same (no shape_id) - no recompilation
        x2 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x2, 0)
        compiled_func(x2)
        self.assertEqual(counter.frame_count, 1)

        # Third call with shape_id="batch" - should recompile
        # (compiled without shape_id, now calling with shape_id)
        x3 = torch.rand(4, 3)
        torch._dynamo.decorators.mark_unbacked(x3, 0, shape_id="batch")
        compiled_func(x3)
        self.assertEqual(counter.frame_count, 2)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_unbacked_exec_fft_reshape_no_dde(self):
        """
        Test that view/reshape operations from with in meta python function
        on tensors with unbacked dimensions do not raise GuardOnDataDependentSymNode
        errors. This tests the unbacked handling in the view decomposition via
        _exec_fft.
        """

        def func(x):
            # FFT triggers _exec_fft which calls reshape with unbacked batch dims
            return torch.fft.fft(x)

        x = torch.rand(4, 8, dtype=torch.complex64)
        torch._dynamo.decorators.mark_unbacked(x, 0)

        torch._dynamo.reset()
        # This should not raise GuardOnDataDependentSymNode
        compiled_func = torch.compile(func, fullgraph=True, backend="eager")
        result = compiled_func(x)
        expected = torch.fft.fft(x)
        self.assertTrue(torch.allclose(result, expected))

    @torch._dynamo.config.patch("capture_scalar_outputs", True)
    def test_inductor_issue_164428(self):
        torch.manual_seed(6804)

        def fuzzed_program(arg_0, arg_1, arg_2, sentinel):
            var_node_4 = arg_0
            var_node_5 = torch.full((7, 1, 32), -1.195053522845565, dtype=torch.float64)
            var_node_3 = torch.div(var_node_4, var_node_5)
            var_node_2 = torch.flatten(var_node_3)
            var_node_8 = torch.full((2,), -0.8316502130341195, dtype=torch.float64)
            var_node_9 = arg_1
            var_node_7 = torch.matmul(
                var_node_8.to(torch.float64), var_node_9.to(torch.float64)
            )
            var_node_10 = arg_2
            var_node_6 = torch.sub(var_node_7, var_node_10)
            var_node_1 = torch.sub(var_node_2, var_node_6)
            var_node_0 = var_node_1.contiguous().view([16, 14])
            result = var_node_0 * sentinel
            if result.is_complex():
                result = result.real
            return result

        sentinel = torch.tensor(1.0, requires_grad=True)

        arg_0 = torch.as_strided(
            torch.randn(7).to(torch.float64), (7, 1, 32), (1, 1, 0)
        )
        arg_1 = torch.as_strided(torch.randn(448).to(torch.float64), (2, 224), (224, 1))
        arg_2 = torch.as_strided(torch.randn(224).to(torch.float64), (224,), (1,))

        args = (arg_0, arg_1, arg_2) + (sentinel,)
        result_original = fuzzed_program(*args)
        compiled_program = torch.compile(fuzzed_program, fullgraph=True, dynamic=True)
        result_compiled = compiled_program(*args)
        self.assertTrue(torch.allclose(result_compiled, result_original))

    def test_inductor_view_issue_172693(self):
        class MinimalBugModule(torch.nn.Module):
            def __init__(self, num_heads: int = 8):
                super().__init__()
                self.num_heads = num_heads

            def forward(self, key):
                reshaped_key = key.reshape(key.shape[0], self.num_heads, -1)
                torch._dynamo.graph_break()
                return reshaped_key + 1

        num_heads = 8

        bug_module = MinimalBugModule(num_heads=num_heads)

        key = torch.empty_strided((2, 16, 9), (288, 9, 1)).normal_()

        with torch.inference_mode():
            eager_result = bug_module(key)
            bug_module.compile(dynamic=True)
            compiled_result = bug_module(key)

        self.assertEqual(compiled_result, eager_result)

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_view_as_complex_unbacked_no_dde(self):
        """
        Test that view_as_complex does not raise GuardOnDataDependentSymNode
        when operating on tensors with unbacked dimensions.
        """

        def func(x):
            return torch.view_as_complex(x)

        # Create a tensor with shape [..., 2] where batch dim is unbacked
        x = torch.rand(4, 8, 2)
        torch._dynamo.decorators.mark_unbacked(x, 0)

        torch._dynamo.reset()
        # This should not raise GuardOnDataDependentSymNode
        compiled_func = torch.compile(func, fullgraph=True, backend="eager")
        result = compiled_func(x)
        expected = torch.view_as_complex(x)
        self.assertTrue(
            torch.allclose(torch.view_as_real(result), torch.view_as_real(expected))
        )

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_view_as_real_backward_unbacked_no_dde(self):
        """
        Test that view_as_real backward (which calls view_as_complex) does not
        raise GuardOnDataDependentSymNode when operating on tensors with
        unbacked dimensions.
        """

        def func(x):
            y = torch.view_as_real(x)
            return y.sum()

        x = torch.rand(4, 8, dtype=torch.complex64, requires_grad=True)
        torch._dynamo.decorators.mark_unbacked(x, 0)

        torch._dynamo.reset()
        # This should not raise GuardOnDataDependentSymNode during backward
        compiled_func = torch.compile(func, fullgraph=True, backend="eager")
        result = compiled_func(x)
        result.backward()

    @skipIfTorchDynamo("mark_unbacked is not traceable")
    def test_mark_unbacked_view_input(self):
        @torch.compile()
        def fn(x):
            return x * 100 + 1

        x = torch.rand(10, 10)
        torch._dynamo.decorators.mark_unbacked(x, 0)
        torch._dynamo.decorators.mark_unbacked(x, 1)
        y = x[0:1, 0:1]  # y is a view, y._base is x
        result = fn(y)
        self.assertEqual(result.shape, (1, 1))


instantiate_parametrized_tests(TestUnbacked)


class TestMaybeFastEvalComparison(TestCase):
    """Tests for _maybe_fast_eval_comparison fast path optimization."""

    def test_sum_of_nonneg_ge_zero(self):
        """Test that sum of non-negative symbols >= 0 returns True."""
        shape_env = ShapeEnv()
        # Create unbacked symbols and constrain them to be positive
        u0 = shape_env.create_unbacked_symint()
        u1 = shape_env.create_unbacked_symint()
        torch._check(u0 >= 0)
        torch._check(u1 >= 0)

        # u0 + u1 >= 0 should be True (both have range [0, inf])
        expr = sympy.GreaterThan(u0.node.expr + u1.node.expr, 0)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertEqual(result, sympy.true)

    def test_single_nonneg_symbol_ge_zero(self):
        """Test that a single non-negative symbol >= 0 returns True."""
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        torch._check(u0 >= 0)

        # u0 >= 0 should be True (u0 has range [0, inf])
        expr = sympy.GreaterThan(u0.node.expr, 0)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertEqual(result, sympy.true)

    def test_zero_le_sum_of_nonneg(self):
        """Test that 0 <= sum of non-negative symbols returns True."""
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        u1 = shape_env.create_unbacked_symint()
        torch._check(u0 >= 0)
        torch._check(u1 >= 0)

        # 0 <= u0 + u1 should be True
        expr = sympy.Le(0, u0.node.expr + u1.node.expr)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertEqual(result, sympy.true)

    def test_nonneg_constant_in_sum(self):
        """Test that sum with non-negative constant >= 0 returns True."""
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        torch._check(u0 >= 0)

        # u0 + 5 >= 0 should be True
        expr = sympy.GreaterThan(u0.node.expr + 5, 0)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertEqual(result, sympy.true)

    def test_negative_constant_returns_none(self):
        """Test that sum with negative constant returns None (undecidable)."""
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        torch._check(u0 >= 0)

        # u0 - 5 >= 0 is undecidable (u0 could be 0, making it -5)
        expr = sympy.GreaterThan(u0.node.expr - 5, 0)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertIsNone(result)

    def test_non_zero_rhs_returns_none(self):
        """Test that comparison with non-zero RHS returns None."""
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        torch._check(u0 >= 0)

        # u0 >= 5 is not handled by this fast path
        expr = sympy.GreaterThan(u0.node.expr, 5)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertIsNone(result)

    def test_unhandled_expr_type_returns_none(self):
        """Test that unhandled expression types return None."""
        shape_env = ShapeEnv()
        u0 = shape_env.create_unbacked_symint()
        torch._check(u0 >= 0)

        # Eq is not handled by the current fast path
        expr = sympy.Eq(u0.node.expr, 0)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertIsNone(result)

    def test_unbacked_symbol_with_default_range_returns_none(self):
        """Test that unbacked symbols with default range [-inf, inf] return None."""
        shape_env = ShapeEnv()
        # Default unbacked symint has range [-int_oo, int_oo]
        u0 = shape_env.create_unbacked_symint()

        # u0 >= 0 is undecidable since u0 could be negative
        expr = sympy.GreaterThan(u0.node.expr, 0)
        result = shape_env._maybe_fast_eval_comparison(expr)
        self.assertIsNone(result)


if __name__ == "__main__":
    run_tests()