torch.fx.replace_pattern can't replace torch.minimum consistently

Viewed 16

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.

0 Answers
Related