active mesh

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -379,28 +379,18 @@ def test_scan_remat_parity(self):
         )
 
         # 3. Run Forward
-        with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-            # For strict parity, we need same weights.
-            # Let's assume initialization with same key gives same weights if structure maps 1:1.
-            # scan/loop structure MIGHT differ if `scan` introduces extra variables? 
-            # Usually strict weight copying is safer.
-            pass
-
-        # Since weight copying might be tricky without access to intricate state, 
-        # let's first check if they run and produce valid shapes.
-        # And if we can, assertions on output closeness IF we force same weights.
-        
         print("Running scan_layers=True...")
-        out_scan = model_scan(**inp_args)["sample"]
-        
-        print("Running scan_layers=False...")
-        # To get same weights, we can try to copy state
-        # nnx.update(model_loop, nnx.state(model_scan)) # Might fail if structure differs
-        out_loop = model_loop(**inp_args)["sample"]
-        
-        print("Running remat_policy='full'...")
-        out_remat = model_remat(**inp_args)["sample"]
-
+        with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+            out_scan = model_scan(**inp_args)["sample"]
+            
+            print("Running scan_layers=False...")
+            # To get same weights, we can try to copy state
+            # nnx.update(model_loop, nnx.state(model_scan)) # Might fail if structure differs
+            out_loop = model_loop(**inp_args)["sample"]
+            
+            print("Running remat_policy='full'...")
+            out_remat = model_remat(**inp_args)["sample"]
+            
         self.assertEqual(out_scan.shape, out_loop.shape)
         self.assertEqual(out_scan.shape, out_remat.shape)