Update

Author: Perseus14

src/maxdiffusion/models/ltx2/logical_sharding_ltx2.py

@@ -60,12 +60,27 @@ class TextEncoderShardingSpecs:
   text_encoder_kernel: Optional[tuple] = None
 
 
+@dataclass
+class TextConnectorShardingSpecs:
+  """Specs for the Text Connector execution."""
+
+  net_0_kernel: tuple = (None, "mlp")
+  net_0_bias: tuple = ("mlp",)
+  net_2_kernel: tuple = ("mlp", None)
+  net_2_bias: tuple = (None,)
+
+
 @dataclass
 class VAEShardingSpecs:
   """Sharding specs for the VAE."""
 
   vae_conv_kernel: Optional[tuple] = None
   force_replication: bool = True
+  # Specs for NNXTimestepEmbedding used in VAE
+  emb_linear_1_kernel: tuple = ("embed", "mlp")
+  emb_linear_1_bias: tuple = ("mlp",)
+  emb_linear_2_kernel: tuple = ("mlp", "embed")
+  emb_linear_2_bias: tuple = ("embed",)
 
 
 # --- Unified Registry for LTX2 ---
@@ -80,6 +95,7 @@ class VAEShardingSpecs:
             use_batched_text_encoder=True,
             text_encoder_kernel=(None, "embed"),
         ),
+        "text_connector": TextConnectorShardingSpecs(),
         "vae": VAEShardingSpecs(vae_conv_kernel=("batch", None, None, None)),
     },
     "trillium": {
@@ -92,6 +108,7 @@ class VAEShardingSpecs:
             use_batched_text_encoder=False,
             text_encoder_kernel=(None, "embed"),
         ),
+        "text_connector": TextConnectorShardingSpecs(),
         "vae": VAEShardingSpecs(vae_conv_kernel=(None, None, None, None)),
     },
 }

src/maxdiffusion/models/ltx2/text_encoders/text_encoders_ltx2.py

@@ -14,7 +14,7 @@
 limitations under the License.
 """
 
-from typing import Tuple, Union, List
+from typing import Tuple, Union, List, Any
 import jax
 import jax.numpy as jnp
 from flax import nnx
@@ -57,6 +57,7 @@ def __init__(
       attention_kernel: str = "flash",
       mesh: jax.sharding.Mesh = None,
       rngs: nnx.Rngs = None,
+      sharding_specs: Optional[Any] = None,
       **kwargs,
   ):
     input_dim = caption_channels * text_proj_in_factor
@@ -82,6 +83,7 @@ def __init__(
         attention_kernel=attention_kernel,
         mesh=mesh,
         rngs=rngs,
+        sharding_specs=sharding_specs,
     )
 
     self.audio_embeddings_connector = Embeddings1DConnector(
@@ -97,6 +99,7 @@ def __init__(
         attention_kernel=attention_kernel,
         mesh=mesh,
         rngs=rngs,
+        sharding_specs=sharding_specs,
     )
 
   def __call__(