MultiHeadAttention only keeps rngs if dropout_rate is positive

Cristian Garcia · Flax Authors · commit 65dfd46515eb · 2025-05-15T23:24:26.000-07:00
PiperOrigin-RevId: 759468680
diff --git a/flax/nnx/nn/attention.py b/flax/nnx/nn/attention.py
@@ -427,7 +427,7 @@ def __init__(
       dot_general_cls=self.out_dot_general_cls,
       rngs=rngs,
     )
-    self.rngs = rngs if keep_rngs else None
+    self.rngs = rngs if keep_rngs and dropout_rate > 0 else None
 
     self.cached_key: nnx.Cache[Array] | None = None
     self.cached_value: nnx.Cache[Array] | None = None

Original file line number	Diff line number	Diff line change
`@@ -427,7 +427,7 @@ def __init__(`
`427`	`427`	`dot_general_cls=self.out_dot_general_cls,`
`428`	`428`	`rngs=rngs,`
`429`	`429`	`)`
`430`		`- self.rngs = rngs if keep_rngs else None`
	`430`	`+ self.rngs = rngs if keep_rngs and dropout_rate > 0 else None`
`431`	`431`
`432`	`432`	`self.cached_key: nnx.Cache[Array] \| None = None`
`433`	`433`	`self.cached_value: nnx.Cache[Array] \| None = None`