explicit sharding to fix scan layers False issue

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1010,17 +1010,12 @@ def scan_fn(carry, block):
             transform_metadata={nnx.PARTITION_NAME: "layers"},
         )(carry, self.transformer_blocks)
       else:
+        activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
+        
         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 = jax.lax.with_sharding_constraint(hidden_states, activation_axis_names)
+            audio_hidden_states = jax.lax.with_sharding_constraint(audio_hidden_states, activation_axis_names)
 
             hidden_states, audio_hidden_states = block(
                 hidden_states=hidden_states,