source sync

PiperOrigin-RevId: 866197055

Author: Google-ML-Automation

src/maxdiffusion/__init__.py

@@ -66,6 +66,7 @@
 
 if is_flax_available():
   from flax import config as flax_config
+
   flax_config.update('flax_always_shard_variable', False)
 
 try:

src/maxdiffusion/trainers/wan_trainer.py

@@ -444,7 +444,7 @@ def loss_fn(params):
     noise = jax.random.normal(key=new_rng, shape=latents.shape, dtype=latents.dtype)
     noisy_latents = scheduler.add_noise(scheduler_state, latents, noise, timesteps)
 
-    with jax.named_scope('forward_pass'):
+    with jax.named_scope("forward_pass"):
       model_pred = model(
           hidden_states=noisy_latents,
           timestep=timesteps,
@@ -453,7 +453,7 @@ def loss_fn(params):
           rngs=nnx.Rngs(dropout_rng),
       )
 
-    with jax.named_scope('loss'):
+    with jax.named_scope("loss"):
       training_target = scheduler.training_target(latents, noise, timesteps)
       training_weight = jnp.expand_dims(scheduler.training_weight(scheduler_state, timesteps), axis=(1, 2, 3, 4))
       loss = (training_target - model_pred) ** 2