Address L6 review feedback: fix weight sharding metadata and missing JAX axis rules context manager

Author: Perseus14

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1312,11 +1312,10 @@ def __call__(
 
       def apply_block(block, context: LTX2BlockContext, mask):
         orig_perturbation_mask = context.perturbation_mask
-        context = dataclasses.replace(context, perturbation_mask=mask)
+        context = context.replace(perturbation_mask=mask)
         with jax.named_scope("Transformer Layer"):
           hidden_states_out, audio_hidden_states_out = block(context)
-        context = dataclasses.replace(
-            context,
+        context = context.replace(
             hidden_states=hidden_states_out.astype(context.hidden_states.dtype),
             audio_hidden_states=audio_hidden_states_out.astype(context.audio_hidden_states.dtype),
             perturbation_mask=orig_perturbation_mask,
@@ -1343,6 +1342,7 @@ def scan_fn_ltx2(carry, block):
               length=self.num_layers,
               in_axes=(nnx.Carry, 0),
               out_axes=(nnx.Carry, 0),
+              transform_metadata={nnx.PARTITION_NAME: "layers"},
           )(carry, self.transformer_blocks)
           hidden_states = final_context.hidden_states
           audio_hidden_states = final_context.audio_hidden_states
@@ -1378,6 +1378,7 @@ def scan_fn_ltx23(carry, block_and_mask):
               length=self.num_layers,
               in_axes=(nnx.Carry, 0),
               out_axes=(nnx.Carry, 0),
+              transform_metadata={nnx.PARTITION_NAME: "layers"},
           )(carry, (self.transformer_blocks, perturbation_mask_per_layer))
           hidden_states = final_context.hidden_states
           audio_hidden_states = final_context.audio_hidden_states

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1614,8 +1614,9 @@ def __call__(
           audio_latents_jax_sharded = audio_latents_jax
 
           if not self.transformer.scan_layers:
-            activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
-            activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed"))
+            with nn_partitioning.axis_rules(self.config.logical_axis_rules):
+              activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
+              activation_axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed"))
             latents_jax_sharded = jax.lax.with_sharding_constraint(latents_jax, activation_axis_names)
             audio_latents_jax_sharded = jax.lax.with_sharding_constraint(audio_latents_jax, activation_axis_names_audio)