Apply safe ruff fixes

Author: Armavica

pymc/_version.py

@@ -29,7 +29,7 @@
 import subprocess
 import sys
 
-from typing import Callable
+from collections.abc import Callable
 
 
 def get_keywords():

pymc/backends/__init__.py

@@ -63,12 +63,11 @@
 
 from collections.abc import Mapping, Sequence
 from copy import copy
-from typing import Optional, Union
+from typing import Optional, TypeAlias, Union
 
 import numpy as np
 
 from pytensor.tensor.variable import TensorVariable
-from typing_extensions import TypeAlias
 
 from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
 from pymc.backends.base import BaseTrace, IBaseTrace
@@ -98,9 +97,9 @@ def _init_trace(
     expected_length: int,
     chain_number: int,
     stats_dtypes: list[dict[str, type]],
-    trace: Optional[BaseTrace],
+    trace: BaseTrace | None,
     model: Model,
-    trace_vars: Optional[list[TensorVariable]] = None,
+    trace_vars: list[TensorVariable] | None = None,
 ) -> BaseTrace:
     """Initializes a trace backend for a chain."""
     strace: BaseTrace
@@ -119,14 +118,14 @@ def _init_trace(
 
 def init_traces(
     *,
-    backend: Optional[TraceOrBackend],
+    backend: TraceOrBackend | None,
     chains: int,
     expected_length: int,
-    step: Union[BlockedStep, CompoundStep],
+    step: BlockedStep | CompoundStep,
     initial_point: Mapping[str, np.ndarray],
     model: Model,
-    trace_vars: Optional[list[TensorVariable]] = None,
-) -> tuple[Optional[RunType], Sequence[IBaseTrace]]:
+    trace_vars: list[TensorVariable] | None = None,
+) -> tuple[RunType | None, Sequence[IBaseTrace]]:
     """Initializes a trace recorder for each chain."""
     if HAS_MCB and isinstance(backend, Backend):
         return init_chain_adapters(

pymc/backends/arviz.py

@@ -163,10 +163,10 @@ def insert(self, k: str, v, idx: int):
 class InferenceDataConverter:
     """Encapsulate InferenceData specific logic."""
 
-    model: Optional[Model] = None
-    posterior_predictive: Optional[Mapping[str, np.ndarray]] = None
-    predictions: Optional[Mapping[str, np.ndarray]] = None
-    prior: Optional[Mapping[str, np.ndarray]] = None
+    model: Model | None = None
+    posterior_predictive: Mapping[str, np.ndarray] | None = None
+    predictions: Mapping[str, np.ndarray] | None = None
+    prior: Mapping[str, np.ndarray] | None = None
 
     def __init__(
         self,
@@ -177,11 +177,11 @@ def __init__(
         log_likelihood=False,
         log_prior=False,
         predictions=None,
-        coords: Optional[CoordSpec] = None,
-        dims: Optional[DimSpec] = None,
-        sample_dims: Optional[list] = None,
+        coords: CoordSpec | None = None,
+        dims: DimSpec | None = None,
+        sample_dims: list | None = None,
         model=None,
-        save_warmup: Optional[bool] = None,
+        save_warmup: bool | None = None,
         include_transformed: bool = False,
     ):
         self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
@@ -466,15 +466,15 @@ def to_inference_data(self):
 def to_inference_data(
     trace: Optional["MultiTrace"] = None,
     *,
-    prior: Optional[Mapping[str, Any]] = None,
-    posterior_predictive: Optional[Mapping[str, Any]] = None,
-    log_likelihood: Union[bool, Iterable[str]] = False,
-    log_prior: Union[bool, Iterable[str]] = False,
-    coords: Optional[CoordSpec] = None,
-    dims: Optional[DimSpec] = None,
-    sample_dims: Optional[list] = None,
+    prior: Mapping[str, Any] | None = None,
+    posterior_predictive: Mapping[str, Any] | None = None,
+    log_likelihood: bool | Iterable[str] = False,
+    log_prior: bool | Iterable[str] = False,
+    coords: CoordSpec | None = None,
+    dims: DimSpec | None = None,
+    sample_dims: list | None = None,
     model: Optional["Model"] = None,
-    save_warmup: Optional[bool] = None,
+    save_warmup: bool | None = None,
     include_transformed: bool = False,
 ) -> InferenceData:
     """Convert pymc data into an InferenceData object.
@@ -543,10 +543,10 @@ def predictions_to_inference_data(
     predictions,
     posterior_trace: Optional["MultiTrace"] = None,
     model: Optional["Model"] = None,
-    coords: Optional[CoordSpec] = None,
-    dims: Optional[DimSpec] = None,
-    sample_dims: Optional[list] = None,
-    idata_orig: Optional[InferenceData] = None,
+    coords: CoordSpec | None = None,
+    dims: DimSpec | None = None,
+    sample_dims: list | None = None,
+    idata_orig: InferenceData | None = None,
     inplace: bool = False,
 ) -> InferenceData:
     """Translate out-of-sample predictions into ``InferenceData``.

pymc/backends/base.py

@@ -25,9 +25,7 @@
 from collections.abc import Mapping, Sequence, Sized
 from typing import (
     Any,
-    Optional,
     TypeVar,
-    Union,
     cast,
 )
 
@@ -53,7 +51,7 @@ class IBaseTrace(ABC, Sized):
     varnames: list[str]
     """Names of tracked variables."""
 
-    sampler_vars: list[dict[str, Union[type, np.dtype]]]
+    sampler_vars: list[dict[str, type | np.dtype]]
     """Sampler stats for each sampler."""
 
     def __len__(self):
@@ -75,7 +73,7 @@ def get_values(self, varname: str, burn=0, thin=1) -> np.ndarray:
         raise NotImplementedError()
 
     def get_sampler_stats(
-        self, stat_name: str, sampler_idx: Optional[int] = None, burn=0, thin=1
+        self, stat_name: str, sampler_idx: int | None = None, burn=0, thin=1
     ) -> np.ndarray:
         """Get sampler statistics from the trace.
 
@@ -219,7 +217,7 @@ def __getitem__(self, idx):
             raise ValueError("Can only index with slice or integer")
 
     def get_sampler_stats(
-        self, stat_name: str, sampler_idx: Optional[int] = None, burn=0, thin=1
+        self, stat_name: str, sampler_idx: int | None = None, burn=0, thin=1
     ) -> np.ndarray:
         """Get sampler statistics from the trace.
 
@@ -443,7 +441,7 @@ def get_values(
         burn: int = 0,
         thin: int = 1,
         combine: bool = True,
-        chains: Optional[Union[int, Sequence[int]]] = None,
+        chains: int | Sequence[int] | None = None,
         squeeze: bool = True,
     ) -> list[np.ndarray]:
         """Get values from traces.
@@ -482,9 +480,9 @@ def get_sampler_stats(
         burn: int = 0,
         thin: int = 1,
         combine: bool = True,
-        chains: Optional[Union[int, Sequence[int]]] = None,
+        chains: int | Sequence[int] | None = None,
         squeeze: bool = True,
-    ) -> Union[list[np.ndarray], np.ndarray]:
+    ) -> list[np.ndarray] | np.ndarray:
         """Get sampler statistics from the trace.
 
         Note: This implementation attempts to squeeze object arrays into a consistent dtype,
@@ -534,7 +532,7 @@ def _slice(self, slice: slice):
         trace._report = self._report._slice(*idxs)
         return trace
 
-    def point(self, idx: int, chain: Optional[int] = None) -> dict[str, np.ndarray]:
+    def point(self, idx: int, chain: int | None = None) -> dict[str, np.ndarray]:
         """Return a dictionary of point values at `idx`.
 
         Parameters

pymc/backends/mcbackend.py

@@ -17,7 +17,7 @@
 import pickle
 
 from collections.abc import Mapping, Sequence
-from typing import Any, Optional, Union, cast
+from typing import Any, cast
 
 import hagelkorn
 import mcbackend as mcb
@@ -144,7 +144,7 @@ def _get_stats(self, fname: str, slc: slice) -> np.ndarray:
         return values
 
     def get_sampler_stats(
-        self, stat_name: str, sampler_idx: Optional[int] = None, burn=0, thin=1
+        self, stat_name: str, sampler_idx: int | None = None, burn=0, thin=1
     ) -> np.ndarray:
         slc = slice(burn, None, thin)
         # When there's just one sampler, default to remove the sampler dimension
@@ -204,7 +204,7 @@ def point(self, idx: int) -> dict[str, np.ndarray]:
 def make_runmeta_and_point_fn(
     *,
     initial_point: Mapping[str, np.ndarray],
-    step: Union[CompoundStep, BlockedStep],
+    step: CompoundStep | BlockedStep,
     model: Model,
 ) -> tuple[mcb.RunMeta, PointFunc]:
     variables, point_fn = get_variables_and_point_fn(model, initial_point)
@@ -254,7 +254,7 @@ def init_chain_adapters(
     backend: mcb.Backend,
     chains: int,
     initial_point: Mapping[str, np.ndarray],
-    step: Union[CompoundStep, BlockedStep],
+    step: CompoundStep | BlockedStep,
     model: Model,
 ) -> tuple[mcb.Run, list[ChainRecordAdapter]]:
     """Create an McBackend metadata description for the MCMC run.

pymc/backends/ndarray.py

@@ -17,7 +17,7 @@
 Store sampling values in memory as a NumPy array.
 """
 
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 
@@ -210,7 +210,7 @@ def _slice_as_ndarray(strace, idx):
 
 
 def point_list_to_multitrace(
-    point_list: list[dict[str, np.ndarray]], model: Optional[Model] = None
+    point_list: list[dict[str, np.ndarray]], model: Model | None = None
 ) -> MultiTrace:
     """transform point list into MultiTrace"""
     _model = modelcontext(model)

pymc/backends/report.py

@@ -16,8 +16,6 @@
 import itertools
 import logging
 
-from typing import Optional
-
 from pymc.stats.convergence import _LEVELS, SamplerWarning
 
 logger = logging.getLogger(__name__)
@@ -44,17 +42,17 @@ def ok(self):
         return all(_LEVELS[warn.level] < _LEVELS["warn"] for warn in self._warnings)
 
     @property
-    def n_tune(self) -> Optional[int]:
+    def n_tune(self) -> int | None:
         """Number of tune iterations - not necessarily kept in trace!"""
         return self._n_tune
 
     @property
-    def n_draws(self) -> Optional[int]:
+    def n_draws(self) -> int | None:
         """Number of draw iterations."""
         return self._n_draws
 
     @property
-    def t_sampling(self) -> Optional[float]:
+    def t_sampling(self) -> float | None:
         """
         Number of seconds that the sampling procedure took.
 

pymc/blocking.py

@@ -20,31 +20,27 @@
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from functools import partial
 from typing import (
     Any,
-    Callable,
     Generic,
     NamedTuple,
-    Optional,
+    TypeAlias,
     TypeVar,
-    Union,
 )
 
 import numpy as np
 
-from typing_extensions import TypeAlias
-
 __all__ = ["DictToArrayBijection"]
 
 
 T = TypeVar("T")
 PointType: TypeAlias = dict[str, np.ndarray]
 StatsDict: TypeAlias = dict[str, Any]
 StatsType: TypeAlias = list[StatsDict]
-StatDtype: TypeAlias = Union[type, np.dtype]
-StatShape: TypeAlias = Optional[Sequence[Optional[int]]]
+StatDtype: TypeAlias = type | np.dtype
+StatShape: TypeAlias = Sequence[int | None] | None
 
 
 # `point_map_info` is a tuple of tuples containing `(name, shape, dtype)` for

pymc/data.py

@@ -18,7 +18,7 @@
 
 from collections.abc import Sequence
 from copy import copy
-from typing import Optional, Union, cast
+from typing import cast
 
 import numpy as np
 import pandas as pd
@@ -203,10 +203,10 @@ def Minibatch(variable: TensorVariable, *variables: TensorVariable, batch_size:
 
 def determine_coords(
     model,
-    value: Union[pd.DataFrame, pd.Series, xr.DataArray],
-    dims: Optional[Sequence[Optional[str]]] = None,
-    coords: Optional[dict[str, Union[Sequence, np.ndarray]]] = None,
-) -> tuple[dict[str, Union[Sequence, np.ndarray]], Sequence[Optional[str]]]:
+    value: pd.DataFrame | pd.Series | xr.DataArray,
+    dims: Sequence[str | None] | None = None,
+    coords: dict[str, Sequence | np.ndarray] | None = None,
+) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None]]:
     """Determines coordinate values from data or the model (via ``dims``)."""
     if coords is None:
         coords = {}
@@ -260,8 +260,8 @@ def ConstantData(
     name: str,
     value,
     *,
-    dims: Optional[Sequence[str]] = None,
-    coords: Optional[dict[str, Union[Sequence, np.ndarray]]] = None,
+    dims: Sequence[str] | None = None,
+    coords: dict[str, Sequence | np.ndarray] | None = None,
     infer_dims_and_coords=False,
     **kwargs,
 ) -> TensorConstant:
@@ -290,8 +290,8 @@ def MutableData(
     name: str,
     value,
     *,
-    dims: Optional[Sequence[str]] = None,
-    coords: Optional[dict[str, Union[Sequence, np.ndarray]]] = None,
+    dims: Sequence[str] | None = None,
+    coords: dict[str, Sequence | np.ndarray] | None = None,
     infer_dims_and_coords=False,
     **kwargs,
 ) -> SharedVariable:
@@ -320,12 +320,12 @@ def Data(
     name: str,
     value,
     *,
-    dims: Optional[Sequence[str]] = None,
-    coords: Optional[dict[str, Union[Sequence, np.ndarray]]] = None,
+    dims: Sequence[str] | None = None,
+    coords: dict[str, Sequence | np.ndarray] | None = None,
     infer_dims_and_coords=False,
-    mutable: Optional[bool] = None,
+    mutable: bool | None = None,
     **kwargs,
-) -> Union[SharedVariable, TensorConstant]:
+) -> SharedVariable | TensorConstant:
     """Data container that registers a data variable with the model.
 
     Depending on the ``mutable`` setting (default: True), the variable