Bumps flax to 0.4.0 (#60)

* Bumps `flax` to `0.4.0`

Updates flax to the most recent version. This currently breaks the
current implementation and way in which rng keys of  dropout is being
handled.

Currently have disabled one of the dropout equivalence tests as I am not
fully aware if there is a method of directly affecting the value of
`next_key` within a treex module.

* [auto] black linting fixes

* fix test_pretrained_flax_module

Co-authored-by: Cristian Garcia <cgarcia.e88@gmail.com>

pyproject.toml

@@ -15,7 +15,7 @@ secondary = true
 
 [tool.poetry.dependencies]
 python = "^3.7"
-flax = "^0.3.4"
+flax = "^0.4.0"
 PyYAML = "^5.4.1"
 rich = "^10.7.0"
 optax = "^0.0.9"

tests/nn/test_dropout.py

@@ -60,7 +60,7 @@ def test_dropout_equivalence(
         # split key same way tx.Dropout does internally
         rng, _ = tx.iter_split(flax_key, 2)
 
-        y_flax = flax_module.apply(variables, x, rng=rng)
+        y_flax = flax_module.apply(variables, x, rngs={"dropout": rng})
         y_treex = treex_module(x)
 
         assert np.allclose(y_flax, y_treex)

tests/nn/test_flax_module.py

@@ -29,10 +29,10 @@ def __call__(self, x, training):
     def test_pretrained_flax_module(self):
         class SomeModule(flax.linen.Module):
             @flax.linen.compact
-            def __call__(self, x, rng, training):
+            def __call__(self, x, training):
                 x = flax.linen.Dense(16)(x)
                 x = flax.linen.BatchNorm()(x, use_running_average=not training)
-                x = flax.linen.Dropout(0.5)(x, deterministic=not training, rng=rng)
+                x = flax.linen.Dropout(0.5)(x, deterministic=not training)
                 x = flax.linen.Conv(16, [3])(x)
 
                 return x
@@ -48,13 +48,12 @@ def __call__(self, x, rng, training):
         variables = flax_module.init(
             {"params": params_key, "dropout": dropout_key},
             x,
-            rng,
             False,
         )
 
         treex_module = tx.FlaxModule(SomeModule(), variables=variables,).init(
             42,
-            inputs=tx.Inputs(x, rng),
+            inputs=tx.Inputs(x),
         )
 
         assert all(
@@ -72,23 +71,26 @@ def __call__(self, x, rng, training):
             )
         )
 
-        y_treex = treex_module(x, rng)
-        y_treex = treex_module(x, rng)
+        flax_next_key = treex_module.next_key.copy()
+        y_treex = treex_module(x)
+        y_treex = treex_module(x)
 
+        rng, next_rng = tx.iter_split(dropout_key)
         y_flax, updates = flax_module.apply(
             variables,
             x,
-            rng,
             training,
             mutable=["batch_stats"],
+            rngs={"dropout": flax_next_key()},
         )
+
         variables = variables.copy(updates)
         y_flax, updates = flax_module.apply(
             variables,
             x,
-            rng,
             training,
             mutable=["batch_stats"],
+            rngs={"dropout": flax_next_key()},
         )
         variables = variables.copy(updates)
 

treex/nn/dropout.py

@@ -89,7 +89,7 @@ def __call__(
             variables,
             x,
             deterministic=not training,
-            rng=rng,
+            rngs={"dropout": rng},
         )
 
         return tp.cast(jnp.ndarray, output)