Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Combine walk_model helper functions in the codebase? #6859

Closed
larryshamalama opened this issue Aug 13, 2023 · 1 comment · Fixed by #6976
Closed

Combine walk_model helper functions in the codebase? #6859

larryshamalama opened this issue Aug 13, 2023 · 1 comment · Fixed by #6976
Labels
maintenance needs info Additional information required

Comments

@larryshamalama
Copy link
Member

Description

We have two walk_models helper functions currently defined and they are very similar: one in pymc/logprob/utils.py (from AePPL's merge) and another in pymc/pytensorf.py.

def walk_model(
graphs: Iterable[TensorVariable],
walk_past_rvs: bool = False,
stop_at_vars: Optional[Set[TensorVariable]] = None,
expand_fn: Callable[[TensorVariable], List[TensorVariable]] = lambda var: [],
) -> Generator[TensorVariable, None, None]:

pymc/pymc/pytensorf.py

Lines 179 to 184 in a4b0581

def walk_model(
graphs: Iterable[TensorVariable],
stop_at_vars: Optional[Set[TensorVariable]] = None,
expand_fn: Callable[[TensorVariable], Iterable[TensorVariable]] = lambda var: [],
) -> Generator[TensorVariable, None, None]:
"""Walk model graphs and yield their nodes.

We can combine the two and remove one due to redundancy. Later on, we could even make it more general, depending on the use case (e.g. in #6834).

CC @ricardoV94 @shreyas3156

@ricardoV94
Copy link
Member

They are subtly different. The walk model in logprob.utils is basically only used for rvs_to_value_vars, which is itself only used in conditional_logp. They are both less flexible/have a specific use that assumes among other things, that we are replacing things in topological order iteratively. It may make more sense to move them to logprob/basic.py and make them internal (leading underscore), to emphasize they are not general

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
maintenance needs info Additional information required
Projects
None yet
Development

Successfully merging a pull request may close this issue.

2 participants