[pt2] Add meta for _add_relu (#140009)

pralay-das · pytorchmergebot · commit f06ee3e5466d · 2024-11-13T06:30:58.000Z
aten._add_relu doesn't have meta function registered, so in dynamic shape case it is throwing an error in dynamo logs: Error: `V1107 11:25:32.344000 140481543555072 torch/_dynamo/symbolic_convert.py:534] [0/1] [__graph_breaks] NotImplementedError: aten::_add_relu.Tensor: attempted to run this operator with Meta tensors, but there was no fake impl or Meta kernel registered. You may have run into this message while using an operator with PT2 compilation APIs (torch.compile/torch.export); in order to use this operator with those APIs you'll need to add a fake impl.` Pull Request resolved: #140009 Approved by: https://github.com/ezyang
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
@@ -3712,6 +3712,14 @@ def meta_relu_(self):
     return self
 
 
+@register_meta(aten._add_relu.Tensor)
+@out_wrapper()
+def meta__add_relu(self, other, alpha=1) -> Tensor:
+    return elementwise_meta(
+        self, other, type_promotion=ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
+    )
+
+
 @register_meta([aten.index_put.default, aten._unsafe_index_put.default])
 def meta_index_put(self, indices, values, accumulate=False):
     return torch.empty_like(self)
