correcting initial point type hinting

Author: Goose

pymc/initial_point.py

@@ -26,6 +26,7 @@
 
 from pymc.logprob.transforms import Transform
 from pymc.pytensorf import (
+    SeedSequenceSeed,
     compile,
     find_rng_nodes,
     replace_rng_nodes,
@@ -67,7 +68,7 @@ def make_initial_point_fns_per_chain(
     overrides: StartDict | Sequence[StartDict | None] | None,
     jitter_rvs: set[TensorVariable] | None = None,
     chains: int,
-) -> list[Callable[[int], PointType]]:
+) -> list[Callable[[SeedSequenceSeed], PointType]]:
     """Create an initial point function for each chain, as defined by initvals.
 
     If a single initval dictionary is passed, the function is replicated for each
@@ -84,7 +85,7 @@ def make_initial_point_fns_per_chain(
 
     Returns
     -------
-    ipfns : list[Callable[[int], dict[str, np.ndarray]]]
+    ipfns : list[Callable[[SeedSequenceSeed], dict[str, np.ndarray]]]
         list of functions that return initial points for each chain.
 
     Raises
@@ -129,7 +130,7 @@ def make_initial_point_fn(
     jitter_rvs: set[TensorVariable] | None = None,
     default_strategy: str = "support_point",
     return_transformed: bool = True,
-) -> Callable[[int], PointType]:
+) -> Callable[[SeedSequenceSeed], PointType]:
     """Create seeded function that computes initial values for all free model variables.
 
     Parameters
@@ -146,7 +147,7 @@ def make_initial_point_fn(
 
     Returns
     -------
-    initial_point_fn : Callable[[int], dict[str, np.ndarray]]
+    initial_point_fn : Callable[[SeedSequenceSeed], dict[str, np.ndarray]]
     """
     sdict_overrides = convert_str_to_rv_dict(model, overrides or {})
     initval_strats = {