Cleaned up ltx2 transformer tests and implementations

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -1,3 +1,7 @@
+"""
+This is a test file used for ensuring numerical parity between pytorch and jax implementation of LTX2.
+This is to be ignored and will not be pushed when commiting to main branch.
+"""
 from typing import Optional, Tuple, Any, Dict, Union
 import jax
 import jax.numpy as jnp
@@ -243,7 +247,7 @@ def __init__(
         rngs=rngs,
         dim=dim,
         dim_out=dim,
-        activation_fn=activation_fn,  # Diffusers uses gelu
+        activation_fn=activation_fn,
         dtype=dtype,
         weights_dtype=weights_dtype,
     )
@@ -633,7 +637,7 @@ def __init__(
         weights_dtype=self.weights_dtype,
     )
 
-    # 3.3. Output Layer Scale/Shift Modulation parameters
+    # 3. Output Layer Scale/Shift Modulation parameters
     param_rng = rngs.params()
     self.scale_shift_table = nnx.Param(
         jax.random.normal(param_rng, (2, inner_dim), dtype=self.weights_dtype) / jnp.sqrt(inner_dim),
@@ -816,7 +820,7 @@ def __call__(
       audio_coords: Optional[jax.Array] = None,
       attention_kwargs: Optional[Dict[str, Any]] = None,
       return_dict: bool = True,
-  ) -> Any:  # Should be AudioVisualModelOutput or Tuple
+  ) -> Any:
     # Determine timestep for audio.
     audio_timestep = audio_timestep if audio_timestep is not None else timestep
 
@@ -974,4 +978,4 @@ def scan_fn(carry, block):
 
     if not return_dict:
       return (output, audio_output)
-    return {"sample": output, "audio_sample": audio_output}  # Placeholder
+    return {"sample": output, "audio_sample": audio_output}