flax
flax copied to clipboard
added `self.param` to `nnx.compat`
Added self.param to nnx.compat.
Also added nnx.compat to API reference.
Things to consider:
- to best mirror
self.paramfrom Linen, we need to callself.rngs.params()implicitly in theself.parammethod, but there's no guarantee that theself.rngsattribute 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 setself.rngs=passed_in_rngand then use that when callingself.param
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.