Merge pull request #3841 from chiamp:dropout2

PiperOrigin-RevId: 626514245

flax/experimental/nnx/nnx/nn/stochastic.py

@@ -75,12 +75,11 @@ def __call__(
     Returns:
       The masked inputs reweighted to preserve mean.
     """
-    rngs = rngs or self.rngs
     deterministic = first_from(
       deterministic,
       self.deterministic,
       error_msg="""No `deterministic` argument was provided to Dropout
-          as either a __call__ argument, class attribute, or nnx.flag.""",
+          as either a __call__ argument or class attribute""",
     )
 
     if (self.rate == 0.0) or deterministic:
@@ -90,10 +89,12 @@ def __call__(
     if self.rate == 1.0:
       return jnp.zeros_like(inputs)
 
-    if rngs is None:
-      raise ValueError(
-        "Dropout needs to generate a random mask but no 'rngs' were provided."
-      )
+    rngs = first_from(
+      rngs,
+      self.rngs,
+      error_msg="""`deterministic` is False, but no `rngs` argument was provided to Dropout
+          as either a __call__ argument or class attribute.""",
+    )
 
     keep_prob = 1.0 - self.rate
     rng = rngs[self.rng_collection]()

flax/experimental/nnx/tests/nn/test_stochastic.py

@@ -14,28 +14,81 @@
 
 
 import jax.numpy as jnp
+import numpy as np
 
 from flax.experimental import nnx
 
+import pytest
+
 
 class TestStochastic:
   def test_dropout_internal_rngs(self):
     n = 0
-    m = nnx.Dropout(rate=0.5, deterministic=False, rngs=nnx.Rngs(dropout=0))
+    m1 = nnx.Dropout(rate=0.5, deterministic=False, rngs=nnx.Rngs(dropout=0))
+    m2 = nnx.Dropout(rate=0.5, deterministic=False)
+    rngs2 = nnx.Rngs(dropout=0)
 
     @nnx.jit
-    def f(m, x):
+    def f(m, x, rngs=None):
       nonlocal n
       n += 1
-      return m(x)
+      return m(x, rngs=rngs)
 
     x = jnp.ones((1, 10))
-    assert m.rngs is not None and m.rngs.dropout.count.value == 0
+    assert m1.rngs is not None and m1.rngs.dropout.count.value == 0
 
-    y = f(m, x)
+    y1 = f(m1, x)
     assert n == 1
-    assert m.rngs.dropout.count.value == 1
+    assert m1.rngs.dropout.count.value == 1
+    y2 = f(m2, x, rngs=rngs2)
+    assert n == 2
+    assert rngs2.dropout.count.value == 1
+    np.testing.assert_allclose(y1, y2)
 
-    y = f(m, x)
-    assert n == 1
-    assert m.rngs.dropout.count.value == 2
+    y1 = f(m1, x)
+    assert m1.rngs.dropout.count.value == 2
+    y2 = f(m2, x, rngs=rngs2)
+    assert rngs2.dropout.count.value == 2
+    np.testing.assert_allclose(y1, y2)
+
+    assert n == 2
+
+  def test_dropout_rng_override(self):
+    m1 = nnx.Dropout(rate=0.5, deterministic=False, rngs=nnx.Rngs(dropout=0))
+    m2 = nnx.Dropout(rate=0.5, deterministic=False, rngs=nnx.Rngs(dropout=1))
+    x = jnp.ones((1, 10))
+
+    y1 = m1(x)
+    y2 = m2(x)
+    with pytest.raises(AssertionError):
+      np.testing.assert_allclose(y1, y2)
+
+    y2 = m2(x, rngs=nnx.Rngs(dropout=0))
+    np.testing.assert_allclose(y1, y2)
+
+  def test_dropout_arg_override(self):
+    m = nnx.Dropout(rate=0.5)
+    x = jnp.ones((1, 10))
+
+    # no deterministic arg provided
+    with pytest.raises(
+      ValueError, match='No `deterministic` argument was provided to Dropout'
+    ):
+      m(x)
+    # deterministic call arg provided
+    m(x, deterministic=True)
+    # deterministic constructor arg provided
+    m.set_attributes(deterministic=True)
+    y = m(x)
+    # both deterministic call and constructor arg provided
+    with pytest.raises(AssertionError):
+      np.testing.assert_allclose(
+        y, m(x, deterministic=False, rngs=nnx.Rngs(dropout=0))
+      )
+    # no rng arg provided
+    m.set_attributes(deterministic=False)
+    with pytest.raises(
+      ValueError,
+      match='`deterministic` is False, but no `rngs` argument was provided to Dropout',
+    ):
+      m(x)