adding parity test file

Author: prishajain1

parity_ltx2_maxdiffusion.py

@@ -50,6 +50,8 @@ def custom_normal(key, shape, dtype=None, **kwargs):
 pipe_module.jax.random.normal = custom_normal
 
 def print_stat(name, t):
+    if hasattr(t, "cpu"):
+        t = t.detach().cpu().float().numpy()
     t_np = np.array(t, dtype=np.float32)
     print(f"[{name}] min: {t_np.min():.5f}, max: {t_np.max():.5f}, mean: {t_np.mean():.5f}, std: {t_np.std():.5f}")