Replace hardcoded PartitionSpec in VACE with logical axis mapping

Author: Perseus14

src/maxdiffusion/models/wan/transformers/transformer_wan_vace.py

@@ -21,7 +21,7 @@
 import jax
 from jax.ad_checkpoint import checkpoint_name
 import jax.numpy as jnp
-from jax.sharding import PartitionSpec
+
 
 from .... import common_types
 from ....configuration_utils import register_to_config
@@ -201,12 +201,12 @@ def __call__(
 
     control_hidden_states = jax.lax.with_sharding_constraint(
         control_hidden_states,
-        PartitionSpec("data", "fsdp", "tensor"),
+        nn.logical_to_mesh_axes(("activation_batch", "activation_length", None)),
     )
     control_hidden_states = checkpoint_name(control_hidden_states, "control_hidden_states")
     encoder_hidden_states = jax.lax.with_sharding_constraint(
         encoder_hidden_states,
-        PartitionSpec("data", "fsdp", None),
+        nn.logical_to_mesh_axes(("activation_batch", "activation_length", None)),
     )
 
     # 1. Self-attention