Update code style for Python 3.9

Author: maresb

pytensor/compile/builders.py

@@ -1,8 +1,9 @@
 """Define new Ops from existing Ops"""
 from collections import OrderedDict
+from collections.abc import Sequence
 from copy import copy
 from functools import partial
-from typing import Dict, List, Optional, Sequence, Tuple, cast
+from typing import Optional, cast
 
 import pytensor.tensor as at
 from pytensor import function
@@ -83,11 +84,11 @@ def local_traverse(out):
 
 def construct_nominal_fgraph(
     inputs: Sequence[Variable], outputs: Sequence[Variable]
-) -> Tuple[
+) -> tuple[
     FunctionGraph,
     Sequence[Variable],
-    Dict[Variable, Variable],
-    Dict[Variable, Variable],
+    dict[Variable, Variable],
+    dict[Variable, Variable],
 ]:
     """Construct an inner-`FunctionGraph` with ordered nominal inputs."""
     dummy_inputs = []
@@ -306,13 +307,13 @@ def _filter_rop_var(inpJ, out):
 
     def __init__(
         self,
-        inputs: List[Variable],
-        outputs: List[Variable],
+        inputs: list[Variable],
+        outputs: list[Variable],
         inline: bool = False,
         lop_overrides: str = "default",
         grad_overrides: str = "default",
         rop_overrides: str = "default",
-        connection_pattern: Optional[List[List[bool]]] = None,
+        connection_pattern: Optional[list[list[bool]]] = None,
         name: Optional[str] = None,
         **kwargs,
     ):

pytensor/compile/function/pfunc.py

@@ -3,8 +3,9 @@
 
 """
 
+from collections.abc import Sequence
 from copy import copy
-from typing import Optional, Sequence, Union, overload
+from typing import Optional, Union, overload
 
 from pytensor.compile.function.types import Function, UnusedInputError, orig_function
 from pytensor.compile.io import In, Out

pytensor/compile/function/types.py

@@ -6,7 +6,7 @@
 import time
 import warnings
 from itertools import chain
-from typing import TYPE_CHECKING, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 
@@ -170,13 +170,13 @@ def validate(self, fgraph):
 
 
 def std_fgraph(
-    input_specs: List[SymbolicInput],
-    output_specs: List[SymbolicOutput],
+    input_specs: list[SymbolicInput],
+    output_specs: list[SymbolicOutput],
     accept_inplace: bool = False,
     fgraph: Optional[FunctionGraph] = None,
-    features: List[Type[Feature]] = [PreserveVariableAttributes],
+    features: list[type[Feature]] = [PreserveVariableAttributes],
     force_clone=False,
-) -> Tuple[FunctionGraph, List[SymbolicOutput]]:
+) -> tuple[FunctionGraph, list[SymbolicOutput]]:
     """Make or set up `FunctionGraph` corresponding to the input specs and the output specs.
 
     Any `SymbolicInput` in the `input_specs`, if its `update` field is not

pytensor/compile/mode.py

@@ -5,7 +5,7 @@
 
 import logging
 import warnings
-from typing import Literal, Optional, Tuple, Union
+from typing import Literal, Optional, Union
 
 from pytensor.compile.function.types import Supervisor
 from pytensor.configdefaults import config
@@ -260,7 +260,7 @@ def apply(self, fgraph):
 # final pass just to make sure
 optdb.register("merge3", MergeOptimizer(), "fast_run", "merge", position=100)
 
-_tags: Union[Tuple[str, str], Tuple]
+_tags: Union[tuple[str, str], tuple]
 
 if config.check_stack_trace in ("raise", "warn", "log"):
     _tags = ("fast_run", "fast_compile")
@@ -548,7 +548,7 @@ def register_mode(name, mode):
     predefined_modes[name] = mode
 
 
-def get_target_language(mode=None) -> Tuple[Literal["py", "c", "numba", "jax"], ...]:
+def get_target_language(mode=None) -> tuple[Literal["py", "c", "numba", "jax"], ...]:
     """Get the compilation target language."""
 
     if mode is None:

pytensor/compile/ops.py

@@ -8,7 +8,6 @@
 import copy
 import pickle
 import warnings
-from typing import Dict, Tuple
 
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
@@ -44,8 +43,8 @@ class ViewOp(COp):
     # Mapping from Type to C code (and version) to use.
     # In the C code, the name of the input variable is %(iname)s,
     # the output variable is %(oname)s.
-    c_code_and_version: Dict = {}
-    __props__: Tuple = ()
+    c_code_and_version: dict = {}
+    __props__: tuple = ()
     _f16_ok: bool = True
 
     def make_node(self, x):
@@ -150,10 +149,10 @@ class DeepCopyOp(COp):
     # Mapping from Type to C code (and version) to use.
     # In the C code, the name of the input variable is %(iname)s,
     # the output variable is %(oname)s.
-    c_code_and_version: Dict = {}
+    c_code_and_version: dict = {}
 
     check_input: bool = False
-    __props__: Tuple = ()
+    __props__: tuple = ()
     _f16_ok: bool = True
 
     def __init__(self):

pytensor/compile/profiling.py

@@ -16,7 +16,7 @@
 import time
 from collections import defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import numpy as np
 
@@ -48,7 +48,7 @@ def extended_open(filename, mode="r"):
 total_graph_rewrite_time: float = 0.0
 total_time_linker: float = 0.0
 
-_atexit_print_list: List["ProfileStats"] = []
+_atexit_print_list: list["ProfileStats"] = []
 _atexit_registered: bool = False
 
 
@@ -234,27 +234,27 @@ def reset(self):
     # Total time spent in Function.vm.__call__
     #
 
-    apply_time: Optional[Dict[Union["FunctionGraph", Variable], float]] = None
+    apply_time: Optional[dict[Union["FunctionGraph", Variable], float]] = None
 
-    apply_callcount: Optional[Dict[Union["FunctionGraph", Variable], int]] = None
+    apply_callcount: Optional[dict[Union["FunctionGraph", Variable], int]] = None
 
-    apply_cimpl: Optional[Dict[Apply, bool]] = None
+    apply_cimpl: Optional[dict[Apply, bool]] = None
     # dict from node -> bool (1 if c, 0 if py)
     #
 
     message: Optional[str] = None
     # pretty string to print in summary, to identify this output
     #
 
-    variable_shape: Dict[Variable, Any] = {}
+    variable_shape: dict[Variable, Any] = {}
     # Variable -> shapes
     #
 
-    variable_strides: Dict[Variable, Any] = {}
+    variable_strides: dict[Variable, Any] = {}
     # Variable -> strides
     #
 
-    variable_offset: Dict[Variable, Any] = {}
+    variable_offset: dict[Variable, Any] = {}
     # Variable -> offset
     #
 
@@ -274,7 +274,7 @@ def reset(self):
 
     linker_node_make_thunks: float = 0.0
 
-    linker_make_thunk_time: Dict = {}
+    linker_make_thunk_time: dict = {}
 
     line_width = config.profiling__output_line_width
 

pytensor/compile/sharedvalue.py

@@ -3,7 +3,7 @@
 import copy
 from contextlib import contextmanager
 from functools import singledispatch
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Optional
 
 from pytensor.graph.basic import Variable
 from pytensor.graph.utils import add_tag_trace
@@ -15,7 +15,7 @@
     from pytensor.graph.type import Type
 
 
-__SHARED_CONTEXT__: Optional[List[Variable]] = None
+__SHARED_CONTEXT__: Optional[list[Variable]] = None
 
 
 @contextmanager

pytensor/configparser.py

@@ -3,6 +3,7 @@
 import shlex
 import sys
 import warnings
+from collections.abc import Sequence
 from configparser import (
     ConfigParser,
     InterpolationError,
@@ -12,7 +13,7 @@
 )
 from functools import wraps
 from io import StringIO
-from typing import Callable, Dict, Optional, Sequence, Union
+from typing import Callable, Optional, Union
 
 from pytensor.utils import hash_from_code
 
@@ -72,7 +73,7 @@ def __init__(self, flags_dict: dict, pytensor_cfg, pytensor_raw_cfg):
         self._flags_dict = flags_dict
         self._pytensor_cfg = pytensor_cfg
         self._pytensor_raw_cfg = pytensor_raw_cfg
-        self._config_var_dict: Dict = {}
+        self._config_var_dict: dict = {}
         super().__init__()
 
     def __str__(self, print_doc=True):

pytensor/gradient.py

@@ -2,21 +2,9 @@
 
 import time
 import warnings
+from collections.abc import Mapping, MutableSequence, Sequence
 from functools import partial, reduce
-from typing import (
-    TYPE_CHECKING,
-    Callable,
-    Dict,
-    List,
-    Literal,
-    Mapping,
-    MutableSequence,
-    Optional,
-    Sequence,
-    Tuple,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Callable, Literal, Optional, TypeVar, Union
 
 import numpy as np
 
@@ -44,7 +32,7 @@
 # TODO: Add `overload` variants
 def as_list_or_tuple(
     use_list: bool, use_tuple: bool, outputs: Union[V, Sequence[V]]
-) -> Union[V, List[V], Tuple[V, ...]]:
+) -> Union[V, list[V], tuple[V, ...]]:
     """Return either a single object or a list/tuple of objects.
 
     If `use_list` is True, `outputs` is returned as a list (if `outputs`
@@ -206,17 +194,17 @@ def Rop(
     """
 
     if not isinstance(wrt, (list, tuple)):
-        _wrt: List[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
+        _wrt: list[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
     else:
         _wrt = [pytensor.tensor.as_tensor_variable(x) for x in wrt]
 
     if not isinstance(eval_points, (list, tuple)):
-        _eval_points: List[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
+        _eval_points: list[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
     else:
         _eval_points = [pytensor.tensor.as_tensor_variable(x) for x in eval_points]
 
     if not isinstance(f, (list, tuple)):
-        _f: List[Variable] = [pytensor.tensor.as_tensor_variable(f)]
+        _f: list[Variable] = [pytensor.tensor.as_tensor_variable(f)]
     else:
         _f = [pytensor.tensor.as_tensor_variable(x) for x in f]
 
@@ -237,7 +225,7 @@ def Rop(
             # Tensor, Sparse have the ndim attribute
             pass
 
-    seen_nodes: Dict[Apply, Sequence[Variable]] = {}
+    seen_nodes: dict[Apply, Sequence[Variable]] = {}
 
     def _traverse(node):
         """TODO: writeme"""
@@ -310,7 +298,7 @@ def _traverse(node):
     for out in _f:
         _traverse(out.owner)
 
-    rval: List[Optional[Variable]] = []
+    rval: list[Optional[Variable]] = []
     for out in _f:
         if out in _wrt:
             rval.append(_eval_points[_wrt.index(out)])
@@ -394,19 +382,19 @@ def Lop(
         If `f` is a list/tuple, then return a list/tuple with the results.
     """
     if not isinstance(eval_points, (list, tuple)):
-        _eval_points: List[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
+        _eval_points: list[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
     else:
         _eval_points = [pytensor.tensor.as_tensor_variable(x) for x in eval_points]
 
     if not isinstance(f, (list, tuple)):
-        _f: List[Variable] = [pytensor.tensor.as_tensor_variable(f)]
+        _f: list[Variable] = [pytensor.tensor.as_tensor_variable(f)]
     else:
         _f = [pytensor.tensor.as_tensor_variable(x) for x in f]
 
     grads = list(_eval_points)
 
     if not isinstance(wrt, (list, tuple)):
-        _wrt: List[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
+        _wrt: list[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
     else:
         _wrt = [pytensor.tensor.as_tensor_variable(x) for x in wrt]
 
@@ -504,7 +492,7 @@ def grad(
         raise TypeError("Cost must be a scalar.")
 
     if not isinstance(wrt, Sequence):
-        _wrt: List[Variable] = [wrt]
+        _wrt: list[Variable] = [wrt]
     else:
         _wrt = list(wrt)
 
@@ -1677,7 +1665,7 @@ def mode_not_slow(mode):
 
 def verify_grad(
     fun: Callable,
-    pt: List[np.ndarray],
+    pt: list[np.ndarray],
     n_tests: int = 2,
     rng: Optional[Union[np.random.Generator, np.random.RandomState]] = None,
     eps: Optional[float] = None,