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

Add option to find_upstream to keep nodes visited even if not found #83

Merged
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions src/qonnx/core/modelwrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -346,16 +346,19 @@ def find_producer(self, tensor_name):
return x
return None

def find_upstream(self, tensor_name, finder_fxn):
def find_upstream(self, tensor_name, finder_fxn, keep_if_not_found=False):
"""Follow the producer chain upstream, calling finder_fxn on each upstream
node until it returns True or there are no nodes left. Returns the list
of nodes visited, or None if finder_fxn did not return True."""
of nodes visited, or None if finder_fxn did not return True. If
keep_if_not_found is specified, returns the list of nodes visited, even
if finder_fxn never returned True, i.e., if the search terminated at an
input or initializer."""
visit_list = []
current_tensor = tensor_name
while True:
current_producer = self.find_producer(current_tensor)
if current_producer is None:
return []
return visit_list if keep_if_not_found else []
else:
found = finder_fxn(current_producer)
visit_list.append(current_producer)
Expand All @@ -364,7 +367,7 @@ def find_upstream(self, tensor_name, finder_fxn):
elif len(current_producer.input) > 0:
current_tensor = current_producer.input[0]
else:
return None
return visit_list if keep_if_not_found else None

def find_consumer(self, tensor_name):
"""Finds and returns the node that consumes the tensor with given name.
Expand Down
Loading