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

[nnx] refactor GraphDef #4630

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[nnx] refactor GraphDef #4630

wants to merge 1 commit into from

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Mar 17, 2025

What does this PR do?

Refactors GraphDef into a flat representation of the graph and refactor flatten and unflatten to use this new representation. This will help speedup traversal when implemented in C++.

class GraphDef(tp.Generic[Node]):
  nodes: list[NodeDef[tp.Any] | VariableDef[tp.Any] | NodeRef[tp.Any]]
  attributes: list[tuple[Key, NodeAttr | ArrayAttr | Static[tp.Any]]]
  num_leaves: int

@cgarciae cgarciae force-pushed the optimize-jit branch 7 times, most recently from 8729bdd to 280d737 Compare March 19, 2025 03:20
@cgarciae cgarciae force-pushed the remove-deprecated branch from a1e3bbf to c77565e Compare March 24, 2025 20:18
@cgarciae cgarciae force-pushed the optimize-jit branch 2 times, most recently from fd90765 to 437cba3 Compare March 25, 2025 05:24
@cgarciae cgarciae force-pushed the remove-deprecated branch from c77565e to 8dd75af Compare March 25, 2025 05:24
@cgarciae cgarciae force-pushed the optimize-jit branch 2 times, most recently from a3a7704 to 33f82b5 Compare March 25, 2025 22:23
@cgarciae cgarciae force-pushed the remove-deprecated branch from 175b832 to ce48e5d Compare March 25, 2025 22:23
Base automatically changed from remove-deprecated to main April 10, 2025 18:37
@cgarciae cgarciae marked this pull request as ready for review April 10, 2025 21:21
@cgarciae cgarciae marked this pull request as draft April 10, 2025 21:25
@cgarciae cgarciae marked this pull request as ready for review April 10, 2025 21:26
@cgarciae cgarciae closed this Apr 10, 2025
@cgarciae cgarciae reopened this Apr 10, 2025
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@cgarciae cgarciae force-pushed the optimize-jit branch 2 times, most recently from ef53cf1 to a196671 Compare April 10, 2025 22:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants