Skip to content

[nnx] mutable array p1 #4715

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
May 7, 2025
Merged

[nnx] mutable array p1 #4715

merged 1 commit into from
May 7, 2025

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Apr 18, 2025

What does this PR do?

  • Adds a FLAX_MUTABLE_ARRAY flag which defaults to false.
  • VariableState will unwrap a MutableArray into Array when accessing .value, a .raw_value property was added to access the content without unwrapping.
  • Rngs now stores the RngStreams inside a streams: dict[str, RngStream] dictionary instead of dynamic attributes, this is a breaking change as the checkpoints will now have an additional streams part component.
  • Object type constructor now has a pytree: 'strict' | 'auto' = 'strict' that defines the style of pytree to create, strict means that the user has to mark of child nodes in __data__, whereas auto uses split and merge to define the flattening and unflattening.

@cgarciae cgarciae force-pushed the nnx-mutable-array-p1 branch 4 times, most recently from e44b2e9 to e3efab4 Compare April 24, 2025 19:02
@cgarciae cgarciae marked this pull request as ready for review April 28, 2025 12:02
@cgarciae cgarciae force-pushed the nnx-mutable-array-p1 branch 6 times, most recently from 92f6a04 to f309de1 Compare May 4, 2025 06:46
@cgarciae cgarciae force-pushed the nnx-mutable-array-p1 branch 6 times, most recently from 460a823 to 433156f Compare May 5, 2025 21:22
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 nnx-mutable-array-p1 branch 4 times, most recently from 5691e11 to 017a1ea Compare May 5, 2025 22:49
@cgarciae cgarciae force-pushed the nnx-mutable-array-p1 branch from 017a1ea to 6ab5f83 Compare May 5, 2025 23:24
@copybara-service copybara-service bot merged commit a7157cf into main May 7, 2025
18 of 19 checks passed
@copybara-service copybara-service bot deleted the nnx-mutable-array-p1 branch May 7, 2025 23:27
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.

1 participant