Fix Graph Break on builtin comparison on NNModule (pytorch#103176)

Fixes pytorch#102338 Pull Request resolved: pytorch#103176 Approved by: https://github.com/anijain2305

Author: msaroufim

test/dynamo/test_modules.py

@@ -1925,6 +1925,33 @@ def fn(x):
  res = opt_fn(x)
  self.assertEqual(ref, res)
 
+ def test_no_graphbreak_builtin_equal(self):
+ class MyModule(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.layer0 = torch.nn.Linear(10, 10)
+ self.layer1 = torch.nn.Linear(10, 10)
+ self.layer2 = torch.nn.Linear(10, 10)
+
+ @property
+ def encoder_layers(self):
+ return [self.layer0, self.layer1, self.layer2]
+
+ def forward(self, x):
+ for layer in self.encoder_layers:
+ output = layer(x)
+ if layer == self.layer0:
+ output = F.relu6(output)
+ else:
+ output = F.relu(output)
+ return output
+
+ x = torch.randn(10, 10)
+
+ m = MyModule()
+
+ opt_m = torch.compile(backend="eager", fullgraph=True)(m)
+
 
 if __name__ == "__main__":
  from torch._dynamo.test_case import run_tests

torch/_dynamo/symbolic_convert.py

@@ -1063,6 +1063,7 @@ def COMPARE_OP(self, inst):
  supported_const_comparison_ops[op](object(), right.value), **options
  )
  )
+
  elif (
  left.is_python_constant()
  and right.is_python_constant()

torch/_dynamo/variables/builtin.py

@@ -1215,6 +1215,7 @@ def _comparison(self, tx, left, right):
  from . import (
  BaseListVariable,
  ConstantVariable,
+ NNModuleVariable,
  TensorVariable,
  UserFunctionVariable,
  )
@@ -1229,6 +1230,13 @@ def _comparison(self, tx, left, right):
  def _unimplemented():
  unimplemented(f"comparison {typestr(left)} {op} {typestr(right)}")
 
+ if (
+ isinstance(left, NNModuleVariable)
+ and isinstance(right, NNModuleVariable)
+ and op in supported_const_comparison_ops
+ ):
+ self.push(ConstantVariable(op(left, right)))
+
  if isinstance(left, UserFunctionVariable):
  if op not in supported_const_comparison_ops.values():
  _unimplemented()