Fix flax.linen.stochastic.Dropout

.pre-commit-config.yaml

@@ -7,6 +7,7 @@
 # Run the hooks on all files with
 # 'pre-commit run --all'
 
+repos:
 - repo: https://github.com/mwouts/jupytext
   rev: v1.13.8
   hooks:

flax/linen/stochastic.py

@@ -60,19 +60,19 @@ def __call__(self, inputs, deterministic: Optional[bool] = None):
     """
     deterministic = merge_param(
         'deterministic', self.deterministic, deterministic)
-    if self.rate == 0.:
+
+    if (self.rate == 0.) or deterministic:
       return inputs
+
     # Prevent gradient NaNs in 1.0 edge-case.
     if self.rate == 1.0:
       return jnp.zeros_like(inputs)
+
     keep_prob = 1. - self.rate
-    if deterministic:
-      return inputs
-    else:
-      rng = self.make_rng(self.rng_collection)
-      broadcast_shape = list(inputs.shape)
-      for dim in self.broadcast_dims:
-        broadcast_shape[dim] = 1
-      mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
-      mask = jnp.broadcast_to(mask, inputs.shape)
-      return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))
+    rng = self.make_rng(self.rng_collection)
+    broadcast_shape = list(inputs.shape)
+    for dim in self.broadcast_dims:
+      broadcast_shape[dim] = 1
+    mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
+    mask = jnp.broadcast_to(mask, inputs.shape)
+    return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))