feat(ltx2): add section comments to all sharding spec classes

Author: Perseus14

src/maxdiffusion/models/ltx2/logical_sharding_ltx2.py

@@ -56,19 +56,21 @@ class LTX2DiTShardingSpecs:
 class TextEncoderShardingSpecs:
   """Specs for the Text Encoder execution."""
 
-  use_batched_text_encoder: bool = False
-  text_encoder_kernel: Optional[tuple] = None
+  use_batched_text_encoder: bool = False  # Flag to batch text encoder execution
+  text_encoder_kernel: Optional[tuple] = None  # Spec for text encoder projection layers
 
 
 @dataclass
 class TextConnectorShardingSpecs:
   """Specs for the Text Connector execution."""
 
+  # --- MLP Specs (NNXSimpleFeedForward) ---
   net_0_kernel: tuple = (None, "mlp")
   net_0_bias: tuple = ("mlp",)
   net_2_kernel: tuple = ("mlp", None)
   net_2_bias: tuple = (None,)
-  # Attention specs (defaulting to trillium/safe defaults)
+
+  # --- Attention Specs (LTX2Attention) ---
   qkv_kernel: tuple = ("embed", "heads")
   out_kernel: tuple = ("heads", "embed")
   out_bias: tuple = ("embed",)
@@ -80,9 +82,11 @@ class TextConnectorShardingSpecs:
 class VAEShardingSpecs:
   """Sharding specs for the VAE."""
 
+  # --- VAE Specific Specs ---
   vae_conv_kernel: Optional[tuple] = None
   force_replication: bool = True
-  # Specs for NNXTimestepEmbedding used in VAE
+
+  # --- Shared Embeddings Specs (NNXTimestepEmbedding) ---
   emb_linear_1_kernel: tuple = ("embed", "mlp")
   emb_linear_1_bias: tuple = ("mlp",)
   emb_linear_2_kernel: tuple = ("mlp", "embed")