Revert mutable kwarg in abstract_init_with_metadata in init checkpoint rule

hemildesai · ashors1 · commit 12adf91f75fc · 2024-05-14T12:00:35.000-07:00
diff --git a/paxml/tasks_lib.py b/paxml/tasks_lib.py
@@ -1787,7 +1787,7 @@ def _apply_init_checkpoint_rule(
       )
     # Initialize with a dummy seed
     var_weight_hparams = ckpt_task.model.abstract_init_with_metadata(
-        inputs_shape_dtype, mutable=DEFAULT_INIT_MUTABLE_LIST)
+        inputs_shape_dtype, extra_mutable_list=DEFAULT_INIT_MUTABLE_LIST)
     ckpt_train_state = ckpt_task.create_train_state_padded_shapes(
         var_weight_hparams)
     train_state_pspecs = ckpt_task.create_train_state_partition_specs(

Original file line number	Diff line number	Diff line change
`@@ -1787,7 +1787,7 @@ def _apply_init_checkpoint_rule(`
`1787`	`1787`	`)`
`1788`	`1788`	`# Initialize with a dummy seed`
`1789`	`1789`	`var_weight_hparams = ckpt_task.model.abstract_init_with_metadata(`
`1790`		`- inputs_shape_dtype, mutable=DEFAULT_INIT_MUTABLE_LIST)`
	`1790`	`+ inputs_shape_dtype, extra_mutable_list=DEFAULT_INIT_MUTABLE_LIST)`
`1791`	`1791`	`ckpt_train_state = ckpt_task.create_train_state_padded_shapes(`
`1792`	`1792`	`var_weight_hparams)`
`1793`	`1793`	`train_state_pspecs = ckpt_task.create_train_state_partition_specs(`