anijain2305/state_change.py

## state_change.py
import torch

class Foo(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("iterations", torch.tensor(0))

    def forward(self, x):
        self.iterations.add_(1)
        return x * self.iterations

x = torch.ones(1, 2, 3)
mod = Foo()
opt_mod = torch.compile(mod, backend="aot_eager", fullgraph=True)

ref = mod(x)
res = opt_mod(x)
assert torch.allclose(ref, res)
# Check again to ensure that mutation is handled correctly
ref = mod(x)
res = opt_mod(x)
assert torch.allclose(ref, res)
	import torch

	class Foo(torch.nn.Module):
	def __init__(self):
	super().__init__()
	self.register_buffer("iterations", torch.tensor(0))

	def forward(self, x):
	self.iterations.add_(1)
	return x * self.iterations

	x = torch.ones(1, 2, 3)
	mod = Foo()
	opt_mod = torch.compile(mod, backend="aot_eager", fullgraph=True)

	ref = mod(x)
	res = opt_mod(x)
	assert torch.allclose(ref, res)
	# Check again to ensure that mutation is handled correctly
	ref = mod(x)
	res = opt_mod(x)
	assert torch.allclose(ref, res)