flax icon indicating copy to clipboard operation
flax copied to clipboard

added `self.param` to `nnx.compat`

Open chiamp opened this issue 1 year ago • 1 comments

Added self.param to nnx.compat. Also added nnx.compat to API reference.

Things to consider:

  • to best mirror self.param from Linen, we need to call self.rngs.params() implicitly in the self.param method, but there's no guarantee that the self.rngs attribute exists; currently I raise an Assertion error if it doesn't exist
  • is there some way to force the user to pass an rng in? For example, if it was mandatory to pass an rng on init (model.init(nnx.Rngs(0), x)), then perhaps we can use that rng passed in to set self.rngs=passed_in_rng and then use that when calling self.param

chiamp avatar Jun 12 '24 23:06 chiamp

Codecov Report

Attention: Patch coverage is 0% with 10 lines in your changes missing coverage. Please review.

Project coverage is 0.00%. Comparing base (31adb00) to head (541603e). Report is 9 commits behind head on main.

Files Patch % Lines
flax/nnx/nnx/compat/module.py 0.00% 10 Missing :warning:
Additional details and impacted files
@@          Coverage Diff          @@
##            main   #3992   +/-   ##
=====================================
  Coverage   0.00%   0.00%           
=====================================
  Files        106     106           
  Lines      13582   13642   +60     
=====================================
- Misses     13582   13642   +60     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

codecov-commenter avatar Jun 13 '24 04:06 codecov-commenter