sharding attempt

prishajain1 · prishajain1 · commit 2e66e3c4ce67 · 2026-02-27T22:08:07.000+05:30
diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py
@@ -57,7 +57,7 @@ def __init__(
         use_bias=True,
         dtype=dtype,
         param_dtype=weights_dtype,
-        kernel_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed", "embed")),
+        kernel_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed")),
         bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
     )
 

Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,7 @@ def __init__(`
`57`	`57`	`use_bias=True,`
`58`	`58`	`dtype=dtype,`
`59`	`59`	`param_dtype=weights_dtype,`
`60`		`- kernel_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed", "embed")),`
	`60`	`+ kernel_init=nnx.with_partitioning(nnx.initializers.zeros, (None, "embed")),`
`61`	`61`	`bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),`
`62`	`62`	`)`
`63`	`63`