needs work

pytorch · Apr 23, 2024 · 6389204 · 6389204
1 parent 9c2ac44
commit 6389204
Show file tree

Hide file tree

Showing 3 changed files with 369 additions and 10 deletions.
diff --git a/torch/_dynamo/variables/higher_order_ops.py b/torch/_dynamo/variables/higher_order_ops.py
@@ -1425,7 +1425,8 @@ def call_function(
         # Proxying user defined functions is not supported.
         inp_args, _ = proxy_args_kwargs(proxied_args, {})
 
-        # Why is this here? Unlike other HOPs, the subgrpah's output for this hop is unrelated
+        # Note:[TemplatedAttention out example value]
+        # Why is this here? Unlike other HOPs, the subgraph's output for this hop is unrelated
         # to what the overall HOP returns, we create the correct output proxy by calling the
         # hop (self.value) with the example values.
         with torch._guards.TracingContext.try_get().fake_mode:

diff --git a/torch/_higher_order_ops/__init__.py b/torch/_higher_order_ops/__init__.py
@@ -1,3 +1,3 @@
 from .cond import cond
 from .while_loop import while_loop
-from .templated_attention import templated_attention
+# from .templated_attention import templated_attention