removed unused arg qk_norm

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -72,7 +72,6 @@ def __init__(
         audio_attention_head_dim: int,
         audio_cross_attention_dim: int,
         activation_fn: str = "gelu",
-        qk_norm: str = "rms_norm_across_heads",
         attention_bias: bool = True,
         attention_out_bias: bool = True,
         norm_elementwise_affine: bool = False,
@@ -435,7 +434,6 @@ def __init__(
         audio_hop_length: int = 160,
         num_layers: int = 48,  # Shared arguments
         activation_fn: str = "gelu",
-        qk_norm: str = "rms_norm_across_heads",
         norm_elementwise_affine: bool = False,
         norm_eps: float = 1e-6,
         caption_channels: int = 3840,
@@ -481,7 +479,6 @@ def __init__(
         self.audio_hop_length = audio_hop_length
         self.num_layers = num_layers
         self.activation_fn = activation_fn
-        self.qk_norm = qk_norm
         self.norm_elementwise_affine = norm_elementwise_affine
         self.norm_eps = norm_eps
         self.caption_channels = caption_channels
@@ -644,7 +641,6 @@ def init_block(rngs):
                 audio_attention_head_dim=self.audio_attention_head_dim,
                 audio_cross_attention_dim=audio_inner_dim,
                 activation_fn=self.activation_fn,
-                qk_norm=self.qk_norm,
                 attention_bias=self.attention_bias,
                 attention_out_bias=self.attention_out_bias,
                 norm_elementwise_affine=self.norm_elementwise_affine,
@@ -676,7 +672,6 @@ def init_block(rngs):
                     audio_attention_head_dim=self.audio_attention_head_dim,
                     audio_cross_attention_dim=audio_inner_dim,
                     activation_fn=self.activation_fn,
-                    qk_norm=self.qk_norm,
                     attention_bias=self.attention_bias,
                     attention_out_bias=self.attention_out_bias,
                     norm_elementwise_affine=self.norm_elementwise_affine,

src/maxdiffusion/tests/ltx2_parity_test.py

@@ -86,6 +86,7 @@ def test_transformer_block_shapes(self):
                 audio_cross_attention_dim=cross_dim,
                 activation_fn="gelu",
                 qk_norm="rms_norm_across_heads",
+                qk_norm="rms_norm_across_heads",
                 mesh=self.mesh,
             )