fix for scan layers = false case

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1012,6 +1012,16 @@ def scan_fn(carry, block):
       else:
         for i, block in enumerate(self.transformer_blocks):
           with jax.named_scope(f"Transformer Block {i}"):
+            graphdef, state = nnx.split(block)
+            
+            def _apply_sharding(x):
+              if hasattr(x, "sharding") and x.sharding is not None:
+                return jax.lax.with_sharding_constraint(x, x.sharding)
+              return x
+              
+            state = jax.tree_util.tree_map(_apply_sharding, state)
+            nnx.update(block, state)
+            
             hidden_states, audio_hidden_states = block(
                 hidden_states=hidden_states,
                 audio_hidden_states=audio_hidden_states,