Wan training: Change checkpoint restore_type to jax.Array

ninatu · martinarroyo · ninatu · commit f08effbb4a49 · 2026-03-31T13:00:36.000Z
- Address review comments: change the restore_type from np.ndarray to jax.Array. This is necessary because JAX sharding has no effect on np.ndarray, using jax.Array ensures that the specified sharding is respected during checkpoint restoration.

Co-authored-by: martinarroyo &lt;martinarroyo@google.com&gt;
diff --git a/src/maxdiffusion/checkpointing/wan_checkpointer_2_1.py b/src/maxdiffusion/checkpointing/wan_checkpointer_2_1.py
@@ -63,7 +63,7 @@ def add_sharding_to_struct(leaf_struct, sharding):
 
     params_restore = ocp.args.PyTreeRestore(
         restore_args=jax.tree.map(
-            lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
+            lambda _: ocp.RestoreArgs(restore_type=jax.Array),
             abstract_train_state_with_sharding,
         )
     )
diff --git a/src/maxdiffusion/checkpointing/wan_checkpointer_i2v_2p1.py b/src/maxdiffusion/checkpointing/wan_checkpointer_i2v_2p1.py
@@ -63,7 +63,7 @@ def add_sharding_to_struct(leaf_struct, sharding):
 
     params_restore = ocp.args.PyTreeRestore(
         restore_args=jax.tree.map(
-            lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
+            lambda _: ocp.RestoreArgs(restore_type=jax.Array),
             abstract_train_state_with_sharding,
         )
     )

Original file line number	Diff line number	Diff line change
`@@ -63,7 +63,7 @@ def add_sharding_to_struct(leaf_struct, sharding):`
`63`	`63`
`64`	`64`	`params_restore = ocp.args.PyTreeRestore(`
`65`	`65`	`restore_args=jax.tree.map(`
`66`		`- lambda _: ocp.RestoreArgs(restore_type=np.ndarray),`
	`66`	`+ lambda _: ocp.RestoreArgs(restore_type=jax.Array),`
`67`	`67`	`abstract_train_state_with_sharding,`
`68`	`68`	`)`
`69`	`69`	`)`