Pass ulysses params to un-scanned LTX2 layers

Perseus14 · Perseus14 · commit 28e54afc6758 · 2026-07-22T08:08:08.000Z
diff --git a/src/maxdiffusion/models/ltx2/transformer_ltx2.py b/src/maxdiffusion/models/ltx2/transformer_ltx2.py
@@ -1090,6 +1090,8 @@ def init_block(rngs):
             flash_block_sizes=flash_block_sizes,
             flash_min_seq_length=self.flash_min_seq_length,
             perturbed_attn=self.perturbed_attn,
+            ulysses_shards=ulysses_shards,
+            ulysses_attention_chunks=ulysses_attention_chunks,
         )
         blocks.append(block)
       self.transformer_blocks = nnx.List(blocks)

Original file line number	Diff line number	Diff line change
`@@ -1090,6 +1090,8 @@ def init_block(rngs):`
`1090`	`1090`	`flash_block_sizes=flash_block_sizes,`
`1091`	`1091`	`flash_min_seq_length=self.flash_min_seq_length,`
`1092`	`1092`	`perturbed_attn=self.perturbed_attn,`
	`1093`	`+ ulysses_shards=ulysses_shards,`
	`1094`	`+ ulysses_attention_chunks=ulysses_attention_chunks,`
`1093`	`1095`	`)`
`1094`	`1096`	`blocks.append(block)`
`1095`	`1097`	`self.transformer_blocks = nnx.List(blocks)`