typing

Author: Disty0

scripts/benchmark_sdnq_inference.py

@@ -1,5 +1,3 @@
-from typing import Optional, Union
-
 import time
 import torch
 from tqdm import tqdm
@@ -35,11 +33,11 @@ def benchmark_linear(name: str, linear: torch.nn.Linear, x: torch.Tensor, steps:
 def main(
     steps: int = 50,
     mnk: int = 8192,
-    dtype: Optional[Union[torch.dtype, str]] = None,
-    device: Optional[str] = None,
-    m: Optional[int] = None,
-    n: Optional[int] = None,
-    k: Optional[int] = None,
+    dtype: torch.dtype | str = None,
+    device: str = None,
+    m: int = None,
+    n: int = None,
+    k: int = None,
 ) -> None:
     if device is None:
         device = "cuda" if torch.cuda.is_available() else "xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else None

scripts/benchmark_sdnq_training.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union, Callable
+from collections.abc import Callable
 
 import time
 import torch
@@ -58,11 +58,11 @@ def benchmark_linear(name: str, linear: Callable, x: torch.Tensor, y: torch.Tens
 def main(
     steps: int = 50,
     mnk: int = 8192,
-    dtype: Optional[Union[torch.dtype, str]] = None,
-    device: Optional[str] = None,
-    m: Optional[int] = None,
-    n: Optional[int] = None,
-    k: Optional[int] = None,
+    dtype: torch.dtype | str = None,
+    device: str = None,
+    m: int = None,
+    n: int = None,
+    k: int = None,
 ) -> None:
     if device is None:
         device = "cuda" if torch.cuda.is_available() else "xpu" if hasattr(torch, "xpu") and torch.xpu.is_available() else None

src/sdnq/dequantizer.py

@@ -1,6 +1,6 @@
 # pylint: disable=protected-access
 
-from typing import Callable
+from collections.abc import Callable
 
 from .common import dtype_dict, conv_types, conv_transpose_types, use_tensorwise_fp8_matmul
 

src/sdnq/forward.py

@@ -1,7 +1,5 @@
 # pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
 
-from typing import List
-
 import torch
 
 from ...common import compile_func, fp_mm_func # noqa: TID252
@@ -18,10 +16,10 @@ def conv_fp16_matmul(
     weight: torch.Tensor,
     scale: torch.FloatTensor,
     result_shape: torch.Size,
-    reversed_padding_repeated_twice: List[int],
+    reversed_padding_repeated_twice: list[int],
     padding_mode: str, conv_type: int,
-    groups: int, stride: List[int],
-    padding: List[int], dilation: List[int],
+    groups: int, stride: list[int],
+    padding: list[int], dilation: list[int],
     bias: torch.FloatTensor = None,
     svd_up: torch.FloatTensor = None,
     svd_down: torch.FloatTensor = None,

src/sdnq/layers/conv/conv_fp16.py

@@ -1,7 +1,5 @@
 # pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
 
-from typing import List
-
 import torch
 
 from ...common import compile_func # noqa: TID252
@@ -17,10 +15,10 @@ def conv_fp8_matmul(
     weight: torch.Tensor,
     scale: torch.FloatTensor,
     result_shape: torch.Size,
-    reversed_padding_repeated_twice: List[int],
+    reversed_padding_repeated_twice: list[int],
     padding_mode: str, conv_type: int,
-    groups: int, stride: List[int],
-    padding: List[int], dilation: List[int],
+    groups: int, stride: list[int],
+    padding: list[int], dilation: list[int],
     bias: torch.FloatTensor = None,
     svd_up: torch.FloatTensor = None,
     svd_down: torch.FloatTensor = None,

src/sdnq/layers/conv/conv_fp8.py

@@ -1,7 +1,5 @@
 # pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
 
-from typing import List
-
 import torch
 
 from ...common import compile_func # noqa: TID252
@@ -18,10 +16,10 @@ def conv_fp8_matmul_tensorwise(
     weight: torch.Tensor,
     scale: torch.FloatTensor,
     result_shape: torch.Size,
-    reversed_padding_repeated_twice: List[int],
+    reversed_padding_repeated_twice: list[int],
     padding_mode: str, conv_type: int,
-    groups: int, stride: List[int],
-    padding: List[int], dilation: List[int],
+    groups: int, stride: list[int],
+    padding: list[int], dilation: list[int],
     bias: torch.FloatTensor = None,
     svd_up: torch.FloatTensor = None,
     svd_down: torch.FloatTensor = None,

src/sdnq/layers/conv/conv_fp8_tensorwise.py

@@ -1,7 +1,5 @@
 # pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
 
-from typing import List
-
 import torch
 
 from ...common import compile_func, int_mm_func # noqa: TID252
@@ -18,10 +16,10 @@ def conv_int8_matmul(
     weight: torch.Tensor,
     scale: torch.FloatTensor,
     result_shape: torch.Size,
-    reversed_padding_repeated_twice: List[int],
+    reversed_padding_repeated_twice: list[int],
     padding_mode: str, conv_type: int,
-    groups: int, stride: List[int],
-    padding: List[int], dilation: List[int],
+    groups: int, stride: list[int],
+    padding: list[int], dilation: list[int],
     bias: torch.FloatTensor = None,
     svd_up: torch.FloatTensor = None,
     svd_down: torch.FloatTensor = None,

src/sdnq/layers/conv/conv_int8.py

@@ -1,7 +1,5 @@
 # pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
 
-from typing import Optional
-
 import torch
 
 
@@ -78,16 +76,16 @@ def quantized_conv_forward(self, input) -> torch.FloatTensor:
     return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias)
 
 
-def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: list[int] = None) -> torch.FloatTensor:
     output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
     return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
 
 
-def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: list[int] = None) -> torch.FloatTensor:
     output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
     return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
 
 
-def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
+def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: list[int] = None) -> torch.FloatTensor:
     output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
     return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)

src/sdnq/layers/conv/forward.py

@@ -1,13 +1,11 @@
 # pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
 
-from typing import Tuple
-
 import torch
 
 from ...common import use_contiguous_mm # noqa: TID252
 
 
-def check_mats(input: torch.Tensor, weight: torch.Tensor, allow_contiguous_mm: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+def check_mats(input: torch.Tensor, weight: torch.Tensor, allow_contiguous_mm: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
     input = input.contiguous()
     if allow_contiguous_mm and use_contiguous_mm:
         weight = weight.contiguous()