change defaults for Dropout and BatchNorm

Changes `Dropout.deterministic` and `BatchNorm.use_running_average` to be None by default, use now has to explicitely provide them by either:

1. Passing them to the constructor e.g:

  self.bn = nnx.BatchNorm(..., use_running_average=False)

2. Passing them to __call__:

  self.dropout(x, deterministic=False)

3. Using `nnx.view` to create a view of the model with specific values:

  train_model = nnx.view(model, detereministic=False, use_running_average=False)

PiperOrigin-RevId: 878147422

Author: Cristian Garcia

src/maxdiffusion/models/attention_flax.py

@@ -1011,7 +1011,7 @@ def __init__(
         ),
     )
 
-    self.drop_out = nnx.Dropout(dropout)
+    self.drop_out = nnx.Dropout(dropout, deterministic=False)
 
     self.norm_q = nnx.data(None)
     self.norm_k = nnx.data(None)

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -237,7 +237,7 @@ def __init__(
     else:
       raise NotImplementedError(f"{activation_fn} is not implemented.")
 
-    self.drop_out = nnx.Dropout(dropout)
+    self.drop_out = nnx.Dropout(dropout, deterministic=False)
     self.proj_out = nnx.Linear(
         rngs=rngs,
         in_features=inner_dim,