Skip to content

Commit

Permalink
add return type overload for sample_posterior_predictive
Browse files Browse the repository at this point in the history
  • Loading branch information
Goose committed Mar 4, 2025
1 parent a386c7a commit 76c3ff0
Showing 1 changed file with 30 additions and 0 deletions.
30 changes: 30 additions & 0 deletions pymc/sampling/forward.py
Original file line number Diff line number Diff line change
Expand Up @@ -493,6 +493,36 @@ def sample_prior_predictive(
return pm.to_inference_data(prior=prior, **ikwargs)


@overload
def sample_posterior_predictive(
trace,
model: Model | None = None,
var_names: list[str] | None = None,
sample_dims: list[str] | None = None,
random_seed: RandomState = None,
progressbar: bool = True,
progressbar_theme: Theme | None = default_progress_theme,
return_inferencedata: Literal[True] = True,
extend_inferencedata: bool = False,
predictions: bool = False,
idata_kwargs: dict | None = None,
compile_kwargs: dict | None = None,
) -> InferenceData: ...
@overload
def sample_posterior_predictive(
trace,
model: Model | None = None,
var_names: list[str] | None = None,
sample_dims: list[str] | None = None,
random_seed: RandomState = None,
progressbar: bool = True,
progressbar_theme: Theme | None = default_progress_theme,
return_inferencedata: Literal[False] = False,
extend_inferencedata: bool = False,
predictions: bool = False,
idata_kwargs: dict | None = None,
compile_kwargs: dict | None = None,
) -> dict[str, np.ndarray]: ...
def sample_posterior_predictive(
trace,
model: Model | None = None,
Expand Down

0 comments on commit 76c3ff0

Please sign in to comment.