insert_deferred_runtime_asserts does not work with modulus

[ezyang]

🐛 Describe the bug

insert_deferred_runtime_asserts operates by symbolic tracing through the PythonReferenceAnalysis. However, the implementation of mod is not traceable:

    @staticmethod
    def mod(x, y): 
        ret = abs(x) % abs(y)
        if x < 0:
            ret *= -1
        return ret

due to the conditional on x.

This is related to @lezcano complaining about how C and Python modulus semantics are not the same.

Versions

main

[ezyang]

Minimal repro:

import torch
import torch._dynamo

torch._dynamo.config.capture_scalar_outputs = True


@torch.compile(fullgraph=True)
def f(x, w):
    y = x.item()
    torch._check(y % 4 == 0)
    return x * w


print(f(torch.tensor(20), torch.randn(3)))

[lezcano]

We need to implement our own C-modulus SymPy class to fix this one.