Revert "[user triton] typing triton_kernel_wrap.py (pytorch#138230)"

This reverts commit 2f61b69.

Reverted pytorch#138230 on behalf of https://github.com/wdvr due to Reverting this, as it started failing tests on main ([comment](pytorch#138230 (comment)))

Author: pytorchmergebot

torch/_dynamo/variables/builder.py

@@ -527,7 +527,18 @@ def _id_dispatch(
 
     def _wrap(self, value):
         # import here to avoid circular dependencies
-        from torch.utils._triton import Autotuner, has_triton_tma, JITFunction
+        from torch.utils._triton import has_triton, has_triton_tma
+
+        if has_triton():
+            from triton.runtime.autotuner import Autotuner
+            from triton.runtime.jit import JITFunction
+        else:
+
+            class JITFunction:
+                pass
+
+            class Autotuner:
+                pass
 
         if has_triton_tma():
             from triton.tools.experimental_descriptor import (

torch/_dynamo/variables/functions.py

@@ -5,18 +5,7 @@
 import inspect
 import itertools
 import types
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    List,
-    Optional,
-    Tuple,
-    TYPE_CHECKING,
-    TypeVar,
-    Union,
-)
-from typing_extensions import Never
+from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING, TypeVar, Union
 
 import torch
 
@@ -47,8 +36,6 @@
 if TYPE_CHECKING:
     from torch._dynamo.symbolic_convert import InstructionTranslator
     from torch._guards import Source
-    from torch._higher_order_ops.triton_kernel_wrap import TritonGridType
-    from torch.utils._triton import TritonKernelType
 
 
 _F = TypeVar("_F", bound=Callable)
@@ -1041,18 +1028,18 @@ def as_python_constant(self):
 
 
 class DynamoTritonHOPifier(TritonHOPifier):
-    def raise_unsupported(self, msg: str) -> Never:
+    def raise_unsupported(self, msg):
         raise Unsupported(msg)
 
-    def is_callable(self, maybe_callable: Any) -> bool:
+    def is_callable(self, maybe_callable):
         return isinstance(
             maybe_callable, (NestedUserFunctionVariable, UserFunctionVariable)
         )
 
-    def get_value(self, val: Any) -> Any:
+    def get_value(self, val):
         return val.value
 
-    def check_grid(self, grid) -> Tuple[torch.fx.proxy.Proxy, ...]:
+    def check_grid(self, grid):
         from .lists import BaseListVariable
 
         if isinstance(grid, BaseListVariable):
@@ -1065,7 +1052,7 @@ def call_grid(self, grid, meta, tx):
         grid = grid.call_function(tx, [meta], {})
         return grid
 
-    def call_HOP(self, variable, grids, combined_args_raw, tx) -> ConstantVariable:
+    def call_HOP(self, variable, grids, combined_args_raw, tx):
         from .constant import ConstantVariable
         from .dicts import ConstDictVariable
 
@@ -1136,10 +1123,6 @@ def call_HOP(self, variable, grids, combined_args_raw, tx) -> ConstantVariable:
 
 
 class TritonKernelVariable(VariableTracker):
-    grid: "TritonGridType"
-    kernel: "TritonKernelType"
-    kernel_idx: Optional[int]
-
     def __init__(self, kernel, kernel_idx, grid, **kwargs) -> None:
         super().__init__(**kwargs)
         dynamo_triton_hopifier_singleton.init_variable(self, kernel, kernel_idx, grid)