Skip to content

Commit 9999e0f

Browse files
committed
Use Mapping instead of Dict in type hints
1 parent 9a3e474 commit 9999e0f

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

pymc_experimental/model_transform/conditioning.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Any, Dict, List, Optional, Sequence, Union
1+
from typing import Any, List, Mapping, Optional, Sequence, Union
22

33
from pymc import Model
44
from pymc.logprob.transforms import RVTransform
@@ -26,7 +26,9 @@
2626
from pymc_experimental.utils.pytensorf import rvs_in_graph
2727

2828

29-
def observe(model: Model, vars_to_observations: Dict[Union["str", TensorVariable], Any]) -> Model:
29+
def observe(
30+
model: Model, vars_to_observations: Mapping[Union["str", TensorVariable], Any]
31+
) -> Model:
3032
"""Convert free RVs or Deterministics to observed RVs.
3133
3234
Parameters
@@ -122,7 +124,9 @@ def replacement_fn(var, inner_replacements):
122124

123125

124126
def do(
125-
model: Model, vars_to_interventions: Dict[Union["str", TensorVariable], Any], prune_vars=False
127+
model: Model,
128+
vars_to_interventions: Mapping[Union["str", TensorVariable], Any],
129+
prune_vars=False,
126130
) -> Model:
127131
"""Replace model variables by intervention variables.
128132
@@ -217,7 +221,7 @@ def do(
217221

218222
def change_value_transforms(
219223
model: Model,
220-
vars_to_transforms: Dict[ModelVariable, Union[RVTransform, None]],
224+
vars_to_transforms: Mapping[ModelVariable, Union[RVTransform, None]],
221225
) -> Model:
222226
"""Change the value variables transforms in the model
223227

0 commit comments

Comments
 (0)