Use PartitionSpec data for input tensor batch sharding

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -242,8 +242,8 @@ def __call__(
       trace["prompt_encoding"] = time.perf_counter() - t0
       max_logging.log(f" -> [TIMING] Prompt Encoding (Qwen3): {trace['prompt_encoding']:.4f} seconds ⏱️")
 
-      # Shard pipeline batch inputs across data_sharding for SPMD multi-host execution
-      data_sharding = jax.sharding.NamedSharding(self.mesh, P(*self._config.data_sharding))
+      # Shard pipeline batch inputs across data axis ("data") for SPMD multi-host execution
+      data_sharding = jax.sharding.NamedSharding(self.mesh, P("data"))
 
       def put_data_on_devices(x, sharding):
         if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable: