[nnx] Rngs and RngStream inherit from GraphNode

[cgarciae]

What does this PR do?

• Rngs and RngStream are now GraphNodes, this means their state can be more easily tracked by transforms.
• Added a RngState(Variable) type.
• RngStream.key and RngStream.count are now RngState variables. This also means that count is now a dynamic property.
• Added a general GraphNode.check_valid_context method that can be used to check that the GraphNode is in the appropriate context.
• Removed Rngs._trace_state and Rngs.is_valid in favor of using GraphNode.check_valid_context.
• nnx.jit no longer has a special rule for Rngs, they are just handled as GraphNodes if present.

other

• Created a ModuleMeta(GraphNodeMeta) and passed all dataclass post-init logic (e.g. calling .setup or setting rngs = None) to ModuleMeta.
• Transform combinator metaclasses now inherit from ModuleMeta.

[codecov-commenter]

Codecov Report

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

Project coverage is 60.20%. Comparing base (514c111) to head (453f38b).

Files	Patch %	Lines
flax/experimental/nnx/nnx/module.py	91.66%	2 Missing ⚠️
flax/experimental/nnx/nnx/variables.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3793      +/-   ##
==========================================
+ Coverage   60.16%   60.20%   +0.03%     
==========================================
  Files         101      101              
  Lines       12842    12859      +17     
==========================================
+ Hits         7727     7742      +15     
- Misses       5115     5117       +2     

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

[chiamp]

Suggested change

	def __init__(self, not_rngs):
	rngs = not_rngs
	self.linear = nnx.Linear(2, 2, rngs=rngs)
	def __init__(self, rngs):
	self.linear = nnx.Linear(2, 2, rngs=rngs)

Small nit: What is the purpose of naming the rng arg not_rngs when it's reassigned to the rngs variable and passed in as an rng key to an nnx.Module anyway?

[chiamp]

Suggested change

	def f(m: Foo, x: jax.Array, not_rngs: nnx.Rngs):
	rngs = not_rngs
	x = m(x, rngs)
	x = m(x, rngs)
	def f(m: Foo, x: jax.Array, rngs: nnx.Rngs):
	x = m(x, rngs)
	x = m(x, rngs)

Same here.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB