Changing data -> None

prishajain1 · prishajain1 · commit c5fb919611ab · 2026-02-26T21:02:45.000+05:30
diff --git a/src/maxdiffusion/models/wan/autoencoder_kl_wan.py b/src/maxdiffusion/models/wan/autoencoder_kl_wan.py
@@ -156,7 +156,7 @@ def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) ->
         shard_width_axis = "context"
 
       x_padded = jax.lax.with_sharding_constraint(
-          x_padded, jax.sharding.PartitionSpec("data", None, shard_axis, shard_width_axis, None)
+          x_padded, jax.sharding.PartitionSpec(None, None, shard_axis, shard_width_axis, None)
       )
 
     out = self.conv(x_padded)

Original file line number	Diff line number	Diff line change
`@@ -156,7 +156,7 @@ def __call__(self, x: jax.Array, cache_x: Optional[jax.Array] = None, idx=-1) ->`
`156`	`156`	`shard_width_axis = "context"`
`157`	`157`
`158`	`158`	`x_padded = jax.lax.with_sharding_constraint(`
`159`		`- x_padded, jax.sharding.PartitionSpec("data", None, shard_axis, shard_width_axis, None)`
	`159`	`+ x_padded, jax.sharding.PartitionSpec(None, None, shard_axis, shard_width_axis, None)`
`160`	`160`	`)`
`161`	`161`
`162`	`162`	`out = self.conv(x_padded)`