Symbolic shapes unable to reason: Ne(Mod(u0*u2 + u1*u2, u0 + u1), 0)

[ezyang]

🐛 Describe the bug

Repro:

import torch
from torch import Tensor
from typing import Optional
import torch._prims_common as utils
import torch._dynamo
        
torch._dynamo.config.capture_scalar_outputs = True


@torch.compile(fullgraph=True)
def f(x):
    a, b, stride = x.tolist()
    torch._check_is_size(a)
    torch._check_is_size(b)
    torch._check_is_size(stride)
    ta = torch.randn(a * stride)
    tb = torch.randn(b * stride)
    r = torch.cat([ta, tb])
    return r.view(a + b, stride)


x = torch.tensor([30, 20, 10])
f(x)

Fails with

torch._dynamo.exc.UserError: Tried to use data-dependent value in the subsequent computation. This can happen when we encounter unbounded dynamic value that is unknown during tracing time.  You will need to explicitly give hint to the compiler. Please take a look at constrain_as_value OR constrain_as_size APIs.  Could not guard on data-dependent expression Ne(Mod(u0*u2 + u1*u2, u0 + u1), 0) (unhinted: Ne(Mod(u0*u2 + u1*u2, u0 + u1), 0)).  (Size-like symbols: u2, u1, u0)

Potential framework code culprit (scroll up for full backtrace):
  File "/data/users/ezyang/b/pytorch/torch/_refs/__init__.py", line 3691, in _reshape_view_helper
    while guard_size_oblivious(accum % length != 0):

For more information, run with TORCH_LOGS="dynamic"

This is obviously false though, just need to be a bit smarter in reasoning.

cc @msaroufim @bdhirsh @anijain2305 @chauhang @lezcano

Versions

main

[lezcano]

For this one, we need to implement a reduction rule similar to the one we have with gcd

[lezcano]

Actually, there's something funny going on here. Is that Mod class ours? Because SymPy can reason this properly:

>>> x = sympy.Symbol("x", integer=True, nonnegative=True)
>>> y= sympy.Symbol("y", integer=True, nonnegative=True)
>>> z = sympy.Symbol("z", integer=True, nonnegative=True)
>>> sympy.Mod(x*z + y*z, x + y)
0

[ezyang]

yup, we have our own

class Mod(sympy.Function):
    """
    We maintain this so that we avoid SymPy correctness issues, such as:
    https://github.com/sympy/sympy/issues/25146
    """

    nargs = (2,)

[ezyang]

@lezcano don't do this one, I'm going to give this as a bootcamp task