cleanup

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -971,7 +971,7 @@ def __call__(
     # 5. Run transformer blocks
     def scan_fn(carry, block):
       hidden_states, audio_hidden_states, rngs_carry = carry
-      with jax.named_scope("Transformer Block i"):
+      with jax.named_scope("Transformer Layer"):
         hidden_states_out, audio_hidden_states_out = block(
             hidden_states=hidden_states,
             audio_hidden_states=audio_hidden_states,
@@ -1010,25 +1010,8 @@ def scan_fn(carry, block):
             transform_metadata={nnx.PARTITION_NAME: "layers"},
         )(carry, self.transformer_blocks)
       else:
-        mlp_rules = nn.logical_to_mesh_axes(("mlp", "tensor"))
-        tensor_rules = nn.logical_to_mesh_axes(("tensor",))
-        
         for i, block in enumerate(self.transformer_blocks):
           with jax.named_scope(f"Transformer Block {i}"):
-            graphdef, state = nnx.split(block)
-            
-            def _apply_weight_sharding(path, x):
-              path_str = "/".join(getattr(p, "name", getattr(p, "key", str(p))) for p in path)
-              if "kernel" in path_str:
-                if "ff" in path_str:
-                  return jax.lax.with_sharding_constraint(x, mlp_rules)
-                elif "attn" in path_str:
-                  return jax.lax.with_sharding_constraint(x, tensor_rules)
-              return x
-              
-            state = jax.tree_util.tree_map_with_path(_apply_weight_sharding, state)
-            nnx.update(block, state)
-            
             hidden_states, audio_hidden_states = block(
                 hidden_states=hidden_states,
                 audio_hidden_states=audio_hidden_states,