Error in documentation on Randomness

[osdf]

I'm going through https://flax.readthedocs.io/en/latest/guides/randomness.html and the shorthands do not work (snippet from the site):

rngs = nnx.Rngs(0, params=1)

# using jax.random
z1 = jax.random.normal(rngs(), (2, 3))
z2 = jax.random.bernoulli(rngs.params(), 0.5, (10,))

# shorthand methods
z1 = rngs.normal((2, 3))                 # generates key from rngs.default
z2 = rngs.params.bernoulli(0.5, (10,)) # generates key from rngs.params

The non-working shorthands also make the very first example on the introductory site
https://flax.readthedocs.io/en/latest/nnx_basics.html not work:

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

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

eventually leads to AttributeError: 'RngStream' object has no attribute 'uniform'.

Is there a way to make this 'shorthand' idea work, or is the only approach to use jax.random.normal etc.?

Thanks!

Flax version: 0.11.1
Jax version: 0.7.0

[vfdev-5]

Hey @osdf , thanks for reporting! I'm trying to figure out the issue you are facing with rngs.
I'm running your snippets as following and do not see any issue:

import jax
import flax
from flax import nnx

print(f"{jax.__version__=}, {flax.__version__=}")


rngs = nnx.Rngs(0, params=1)

# using jax.random
z1 = jax.random.normal(rngs(), (2, 3))
z2 = jax.random.bernoulli(rngs.params(), 0.5, (10,))
print(z1, z2)

# shorthand methods
z1 = rngs.normal((2, 3))                 # generates key from rngs.default
z2 = rngs.params.bernoulli(0.5, (10,)) # generates key from rngs.params
print(z1, z2)

Output:

jax.__version__='0.7.0', flax.__version__='0.11.1'
[[ 1.0040143  -0.9063372  -0.7481722 ]
 [-1.1713669  -0.8712328   0.58883816]] [False  True False  True  True False False  True False  True]
[[-2.4424558  -2.0356805   0.20554423]
 [-0.3535502  -0.76197404 -1.1785518 ]] [ True False  True  True False False  True  True  True False]

and

import jax
import jax.numpy as jnp
import flax
from flax import nnx

print(f"{jax.__version__=}, {flax.__version__=}")


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

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

model = Linear(2, 5, rngs=nnx.Rngs(params=0))
y = model(x=jnp.ones((1, 2)))

print(y)
nnx.display(model)

Output:

jax.__version__='0.7.0', flax.__version__='0.11.1'
[[1.5643291  0.94782424 0.37971854 1.0724319  0.22112393]]
Linear( # Param: 15 (60 B)
  w=Param( # 10 (40 B)
    value=Array(shape=(2, 5), dtype=dtype('float32'))
  ),
  b=Param( # 5 (20 B)
    value=Array(shape=(5,), dtype=dtype('float32'))
  ),
  din=2,
  dout=5
)

Can you provide more details on what does not work? Thanks!

[osdf]

Thanks @vfdev-5 for your reply. I'm seeing the following error message for your first snippet:

jax.__version__='0.7.0', flax.__version__='0.11.1'
[[ 1.0040143 -0.9063372 -0.7481722]
 [-1.1713669 -0.8712328  0.5888381]] [False  True False  True  True False False  True False  True]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[1], line 16
     13 print(z1, z2)
     15 # shorthand methods
---> 16 z1 = rngs.normal((2, 3))                 # generates key from rngs.default
     17 z2 = rngs.params.bernoulli(0.5, (10,)) # generates key from rngs.params
     18 print(z1, z2)

TypeError: RngStream.__call__() takes 1 positional argument but 2 were given

I wonder what I'm doing wrong?

[vfdev-5]

Interesting, what is your OS, python version, how did you install jax, flax? Can you retry with a clean python env?
I'll retry the same from my side to see if I can repro the problem.

[osdf]

Python version: 3.12.11 (main, Jul 23 2025, 00:18:05) [Clang 20.1.4 ]
Operating system: Darwin 24.6.0

Originally I installed jax/flax using uv. I did now a pip-based installation and found that if I have flax 0.11.2, rngs.normal((2,4)) works, it fails with flax 0.11.1.

[vfdev-5]

@osdf I could reproduce the issue and figure out what happens, FLAX docs shows the latest version of the docs pointing to what is on main branch. Rng shortcuts were introduced two weeks ago in this PR: #4899 and this updated the docs, but the code was not yet in the stable release v0.11.1 (and appeared in v0.11.2) and that's why you saw the error.

I close the issue as resolved, feel free to reopen if needed more help on that or please open another issue for other topics!