I am trying to rewrite PyTorch modules to eliminate some operations by replacing them by others using torch.fx.
The torch.fx.replace_pattern function seems promising and works well for simple examples.
However, I stumbled upon a somewhat complex pattern that it can't match. I am wondering whether that is intended or not --- I couldn't find anything about it in the documentation.
For demo purposes, I am trying to replace torch.minimum by torch.maximum:
import torch
from torch import nn, fx
class TestModule(nn.Module):
def forward(self, x, y):
y = torch.minimum(x, y) + torch.minimum(x, y)
return torch.minimum(y, x)
def pattern(x, y):
return torch.minimum(x, y)
def replacement(x, y):
return torch.maximum(x, y)
module = TestModule()
traced = fx.symbolic_trace(module)
fx.replace_pattern(traced, pattern, replacement)
print(traced)
assert "minimum" not in str(traced)
The output shows that torch.fx.replace_pattern was able to replace the later calls to torch.minimum, but it didn't replace the first one.
I was able to replace all occurrences using a custom torch.fx.Transformer, but I am wondering whether I have to expect more of the above when using replace_pattern.