Merge pull request #55 from TCLResearchEurope/add-rsqrt

add torch.rsqrt operation support

Author: michal-lopuszynski-tcl

torch_dag/core/unstructured_to_structured.py

@@ -383,6 +383,9 @@ def convert_node(self, node: torch.fx.node.Node, modules_dict, state_dict):
             assert len(node.args) == 1
             return structured_modules.MeanModule(**node.kwargs)
 
+        elif node.target == torch.rsqrt:
+            return structured_modules.RsqrtModule()
+
         elif node.target == torch.sum:
             raise NotImplementedError
 

torch_dag/structured_modules.py

@@ -1588,6 +1588,14 @@ def __init__(self, pow: Union[float, int]):
     def forward(self, inputs: torch.Tensor):
         return torch.pow(inputs, self.pow)
 
+@register_notrace_module
+class RsqrtModule(torch.nn.Module):
+    # wraps torch.rsqrt
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x) -> torch.Tensor:
+        return torch.rsqrt(x)
 
 @register_notrace_module
 class UnsqueezeModule(torch.nn.Module):