fix

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -247,9 +247,19 @@ def __init__(
 
       # Initialize video processor
       self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio)
-      
       self.tokenizer_max_length = getattr(self.tokenizer, "model_max_length", 1024)
 
+  @staticmethod
+  def _init_dummy_shape(node):
+      if isinstance(node, jax.ShapeDtypeStruct):
+          if jax.dtypes.issubdtype(node.dtype, jax.dtypes.prng_key):
+              dummy_key = jax.random.key(0)
+              if node.shape == ():
+                  return dummy_key
+              return jax.random.split(dummy_key, node.shape[0])
+          return jnp.zeros(node.shape, dtype=node.dtype)
+      return node
+
   @classmethod
   def load_tokenizer(cls, config: HyperParameters):
       max_logging.log("Loading Gemma Tokenizer...")
@@ -286,7 +296,8 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
 
       p_model_factory = partial(create_model, config=config)
       connectors = nnx.eval_shape(p_model_factory, rngs=rngs)
-      graphdef, state = nnx.split(connectors, nnx.Param)
+      graphdef, state, rest_of_state = nnx.split(connectors, nnx.Param, ...)
+      rest_of_state = jax.tree_util.tree_map(cls._init_dummy_shape, rest_of_state)
 
       logical_state_spec = nnx.get_partition_spec(state)
       logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
@@ -307,7 +318,7 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
              state[path].value = jax.device_put(val)
 
       state = nnx.from_flat_state(state)
-      connectors = nnx.merge(graphdef, state)
+      connectors = nnx.merge(graphdef, state, rest_of_state)
       return connectors
 
   @classmethod
@@ -326,7 +337,8 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
 
       p_model_factory = partial(create_model, config=config)
       vae = nnx.eval_shape(p_model_factory, rngs=rngs)
-      graphdef, state = nnx.split(vae, nnx.Param)
+      graphdef, state, rest_of_state = nnx.split(vae, nnx.Param, ...)
+      rest_of_state = jax.tree_util.tree_map(cls._init_dummy_shape, rest_of_state)
 
       logical_state_spec = nnx.get_partition_spec(state)
       logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
@@ -353,7 +365,7 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
              state[path].value = jax.device_put(val)
 
       state = nnx.from_flat_state(state)
-      vae = nnx.merge(graphdef, state)
+      vae = nnx.merge(graphdef, state, rest_of_state)
       return vae
 
   @classmethod
@@ -372,7 +384,8 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
 
       p_model_factory = partial(create_model, config=config)
       audio_vae = nnx.eval_shape(p_model_factory, rngs=rngs)
-      graphdef, state = nnx.split(audio_vae, nnx.Param)
+      graphdef, state, rest_of_state = nnx.split(audio_vae, nnx.Param, ...)
+      rest_of_state = jax.tree_util.tree_map(cls._init_dummy_shape, rest_of_state)
 
       logical_state_spec = nnx.get_partition_spec(state)
       logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
@@ -399,7 +412,7 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
              state[path].value = jax.device_put(val)
 
       state = nnx.from_flat_state(state)
-      audio_vae = nnx.merge(graphdef, state)
+      audio_vae = nnx.merge(graphdef, state, rest_of_state)
       return audio_vae
 
   @classmethod
@@ -439,7 +452,8 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
 
       p_model_factory = partial(create_model, config=config)
       vocoder = nnx.eval_shape(p_model_factory, rngs=rngs)
-      graphdef, state = nnx.split(vocoder, nnx.Param)
+      graphdef, state, rest_of_state = nnx.split(vocoder, nnx.Param, ...)
+      rest_of_state = jax.tree_util.tree_map(cls._init_dummy_shape, rest_of_state)
 
       logical_state_spec = nnx.get_partition_spec(state)
       logical_state_sharding = nn.logical_to_mesh_sharding(logical_state_spec, mesh, config.logical_axis_rules)
@@ -460,7 +474,7 @@ def create_model(rngs: nnx.Rngs, config: HyperParameters):
              state[path].value = jax.device_put(val)
 
       state = nnx.from_flat_state(state)
-      vocoder = nnx.merge(graphdef, state)
+      vocoder = nnx.merge(graphdef, state, rest_of_state)
       return vocoder
 
   @classmethod

test_keys.py

@@ -0,0 +1,24 @@
+import jax
+import jax.numpy as jnp
+from flax import nnx
+
+def fix_struct(x):
+    if isinstance(x, jax.ShapeDtypeStruct):
+        if jax.dtypes.issubdtype(x.dtype, jax.dtypes.prng_key):
+            key = jax.random.key(0)
+            if x.shape == ():
+                return key
+            else:
+                return jax.random.split(key, x.shape[0])
+        else:
+            return jnp.zeros(x.shape, x.dtype)
+    return x
+
+struct_key = jax.ShapeDtypeStruct((5,), jax.dtypes.prng_key)
+struct_count = jax.ShapeDtypeStruct((5,), jnp.uint32)
+
+fixed_key = fix_struct(struct_key)
+fixed_count = fix_struct(struct_count)
+
+print("Key dtype:", fixed_key.dtype, "shape:", fixed_key.shape)
+print("Count dtype:", fixed_count.dtype, "shape:", fixed_count.shape)