Skip to content

[nnx] explicit Variables#3720

Merged
copybara-service[bot] merged 1 commit into
mainfrom
nnx-explicit-variables
Mar 1, 2024
Merged

[nnx] explicit Variables#3720
copybara-service[bot] merged 1 commit into
mainfrom
nnx-explicit-variables

Conversation

@cgarciae

@cgarciae cgarciae commented Feb 26, 2024 •
edited
Loading

Copy link
Copy Markdown
Collaborator

What does this PR do?

  • The Variable.value attribute which contained the unprocessed values is now named .raw_value.
  • Replaces Variable.get_value and Variable.set_value with a .value property.
  • Module getattr and setattr no longer extract the inner value of Variables.
  • Removes Module.variables and State.variables helpers.

Basic example now looks like this:

class Linear(nnx.Module):
  def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
    key = rngs.params()
    self.w: Param[Array] = nnx.Param(jax.random.uniform(key, (din, dout)))
    self.b: Param[Array] = nnx.Param(jnp.zeros((dout,)))

  def __call__(self, x: jax.Array):
    return x @ self.w.value + self.b.value

@review-notebook-app

review-notebook-app Bot commented Feb 26, 2024

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

@cgarciae cgarciae force-pushed the nnx-explicit-variables branch from 4fb9e11 to 5f5dd38 Compare February 26, 2024 18:24
@codecov-commenter

codecov-commenter commented Feb 26, 2024 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

Attention: Patch coverage is 95.66474% with 15 lines in your changes are missing coverage. Please review.

Project coverage is 58.85%. Comparing base (acba0bf) to head (0d347d0).

Files Patch % Lines
flax/experimental/nnx/nnx/variables.py 87.75% 6 Missing ⚠️
flax/experimental/nnx/nnx/spmd.py 57.14% 3 Missing ⚠️
flax/experimental/nnx/nnx/state.py 70.00% 3 Missing ⚠️
flax/experimental/nnx/nnx/nn/linear.py 83.33% 2 Missing ⚠️
flax/experimental/nnx/nnx/nn/normalization.py 66.66% 1 Missing ⚠️
Additional details and impacted files 
@@            Coverage Diff             @@
##             main    #3720      +/-   ##
==========================================
- Coverage   59.05%   58.85%   -0.21%     
==========================================
  Files         103      103              
  Lines       12440    12318     -122     
==========================================
- Hits         7347     7250      -97     
+ Misses       5093     5068      -25

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@cgarciae cgarciae force-pushed the nnx-explicit-variables branch 2 times, most recently from 37450b8 to 6ad6c96 Compare February 29, 2024 22:03
@cgarciae cgarciae force-pushed the nnx-explicit-variables branch 4 times, most recently from 4d5de7f to 4997d59 Compare March 1, 2024 14:35
@cgarciae cgarciae force-pushed the nnx-explicit-variables branch 2 times, most recently from d1a67cf to 0d347d0 Compare March 1, 2024 15:30
@copybara-service copybara-service Bot merged commit 7d6de00 into main Mar 1, 2024
@copybara-service copybara-service Bot deleted the nnx-explicit-variables branch March 1, 2024 22:11
@cgarciae cgarciae mentioned this pull request Mar 7, 2024
Open
This was referenced Mar 7, 2024
Merged
Merged
@NeilGirdhar

NeilGirdhar commented Mar 8, 2024

Copy link
Copy Markdown
Contributor

Woohoo!!! Awesome!!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@levskaya levskaya levskaya approved these changes

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@cgarciae @codecov-commenter @NeilGirdhar @levskaya