Skip to content

[nnx] refactor GraphDef #4630

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

Merged
merged 1 commit into from
Apr 15, 2025
Merged

[nnx] refactor GraphDef #4630

merged 1 commit into from
Apr 15, 2025

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 3 times, most recently from a196671 to 2a5df33 Compare April 10, 2025 23:04
@cgarciae cgarciae closed this Apr 10, 2025
@cgarciae cgarciae reopened this Apr 10, 2025
@cgarciae cgarciae force-pushed the optimize-jit branch 3 times, most recently from 6a9f0bf to d7a2a6f Compare April 14, 2025 00:25
@cgarciae cgarciae changed the title [nnx] refactor GraphDef [nnx] refactor GraphDef + remove PPO example Apr 14, 2025
@cgarciae cgarciae changed the title [nnx] refactor GraphDef + remove PPO example [nnx] refactor GraphDef Apr 14, 2025
@copybara-service copybara-service bot merged commit 1238a61 into main Apr 15, 2025
12 of 18 checks passed
@copybara-service copybara-service bot deleted the optimize-jit branch April 15, 2025 17:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants