Fix tuple-related type hints

Author: michaelosthege

pytensor/graph/op.py

@@ -41,7 +41,7 @@
 ComputeMapType = Dict[Variable, List[bool]]
 InputStorageType = List[StorageCellType]
 OutputStorageType = List[StorageCellType]
-ParamsInputType = Optional[Tuple[Any]]
+ParamsInputType = Optional[Tuple[Any, ...]]
 PerformMethodType = Callable[
     [Apply, List[Any], OutputStorageType, ParamsInputType], None
 ]

pytensor/link/c/interface.py

@@ -1,6 +1,6 @@
 import warnings
 from abc import abstractmethod
-from typing import Callable, Dict, List, Tuple, Union
+from typing import Callable, Dict, List, Tuple
 
 from pytensor.graph.basic import Apply, Constant
 from pytensor.graph.utils import MethodNotDefined
@@ -149,7 +149,7 @@ def c_init_code(self, **kwargs) -> List[str]:
         """Return a list of code snippets to be inserted in module initialization."""
         return []
 
-    def c_code_cache_version(self) -> Union[Tuple[int, ...], Tuple]:
+    def c_code_cache_version(self) -> Tuple[int, ...]:
         """Return a tuple of integers indicating the version of this `Op`.
 
         An empty tuple indicates an "unversioned" `Op` that will not be cached
@@ -566,7 +566,7 @@ def c_cleanup(self, name: str, sub: Dict[str, str]) -> str:
         """
         return ""
 
-    def c_code_cache_version(self) -> Union[Tuple, Tuple[int]]:
+    def c_code_cache_version(self) -> Tuple[int, ...]:
         """Return a tuple of integers indicating the version of this type.
 
         An empty tuple indicates an "unversioned" type that will not

pytensor/link/c/op.py

@@ -240,7 +240,7 @@ def lquote_macro(txt: str) -> str:
     return "\n".join(res)
 
 
-def get_sub_macros(sub: Dict[str, str]) -> Union[Tuple[str], Tuple[str, str]]:
+def get_sub_macros(sub: Dict[str, str]) -> Tuple[str, str]:
     define_macros = []
     undef_macros = []
     define_macros.append(f"#define FAIL {lquote_macro(sub['fail'])}")
@@ -533,7 +533,7 @@ def format_c_function_args(self, inp: List[str], out: List[str]) -> str:
 
     def get_c_macros(
         self, node: Apply, name: str, check_input: Optional[bool] = None
-    ) -> Union[Tuple[str], Tuple[str, str]]:
+    ) -> Tuple[str, str]:
         "Construct a pair of C ``#define`` and ``#undef`` code strings."
         define_template = "#define %s %s"
         undef_template = "#undef %s"

pytensor/link/numba/dispatch/scan.py

@@ -123,7 +123,7 @@ def add_inner_in_expr(
         # These outer-inputs are indexed without offsets or storage wrap-around
         add_inner_in_expr(outer_in_name, 0, None)
 
-    inner_in_names_to_input_taps: Dict[str, Tuple[int]] = dict(
+    inner_in_names_to_input_taps: Dict[str, Tuple[int, ...]] = dict(
         zip(
             outer_in_mit_mot_names + outer_in_mit_sot_names + outer_in_sit_sot_names,
             op.info.mit_mot_in_slices
@@ -157,7 +157,7 @@ def add_inner_in_expr(
     # storage array like a circular buffer, and that's why we need to track the
     # storage size along with the taps length/indexing offset.
     def add_output_storage_post_proc_stmt(
-        outer_in_name: str, tap_sizes: Tuple[int], storage_size: str
+        outer_in_name: str, tap_sizes: Tuple[int, ...], storage_size: str
     ):
 
         tap_size = max(tap_sizes)

pytensor/raise_op.py

@@ -1,7 +1,6 @@
 """Symbolic Op for raising an exception."""
 
 from textwrap import indent
-from typing import Tuple
 
 import numpy as np
 
@@ -63,7 +62,7 @@ def __eq__(self, other):
     def __hash__(self):
         return hash((self.msg, self.exc_type))
 
-    def make_node(self, value: Variable, *conds: Tuple[Variable]):
+    def make_node(self, value: Variable, *conds: Variable):
         """
 
         Parameters

pytensor/sandbox/multinomial.py

@@ -1,5 +1,5 @@
 import copy
-from typing import Tuple, Union
+from typing import Tuple
 
 import numpy as np
 
@@ -18,7 +18,7 @@ class MultinomialFromUniform(COp):
     TODO : need description for parameter 'odtype'
     """
 
-    __props__: Union[Tuple[str], Tuple[str, str]] = ("odtype",)
+    __props__: Tuple[str, ...] = ("odtype",)
 
     def __init__(self, odtype):
         self.odtype = odtype

pytensor/scalar/basic.py

@@ -3998,7 +3998,7 @@ class Composite(ScalarOp, HasInnerGraph):
 
     """
 
-    init_param: Union[Tuple[str, str], Tuple[str]] = ("inputs", "outputs")
+    init_param: Tuple[str, ...] = ("inputs", "outputs")
 
     def __init__(self, inputs, outputs):
         # We need to clone the graph as sometimes its nodes already

pytensor/tensor/blas.py

@@ -137,7 +137,7 @@
 except ImportError:
     pass
 
-from typing import Tuple, Union
+from typing import Tuple
 
 import pytensor.scalar
 from pytensor.compile.mode import optdb
@@ -522,7 +522,7 @@ class GemmRelated(COp):
 
     """
 
-    __props__: Union[Tuple, Tuple[str]] = ()
+    __props__: Tuple[str, ...] = ()
 
     def c_support_code(self, **kwargs):
         # return cblas_header_text()

pytensor/tensor/extra_ops.py

@@ -1763,7 +1763,7 @@ def linspace(start, end, steps):
 
 
 def broadcast_to(
-    x: TensorVariable, shape: Union[TensorVariable, Tuple[Variable]]
+    x: TensorVariable, shape: Union[TensorVariable, Tuple[Variable, ...]]
 ) -> TensorVariable:
     """Broadcast an array to a new shape.
 

pytensor/tensor/nlinalg.py

@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Tuple, Union
+from typing import Tuple
 
 import numpy as np
 
@@ -238,7 +238,7 @@ class Eig(Op):
     """
 
     _numop = staticmethod(np.linalg.eig)
-    __props__: Union[Tuple, Tuple[str]] = ()
+    __props__: Tuple[str, ...] = ()
 
     def make_node(self, x):
         x = as_tensor_variable(x)