fix(nnx): support non-indexed opt_state paths in gpt3 paxml conversion

Author: hsuan-lun-chiang

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -213,6 +213,7 @@ def init_state_fn():
     state_map = {
         "['optimizer']['step'].value": ("step", None),
         "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None),
+        "['optimizer']['opt_state']['count'].value": ("opt_states_0.no_prefix_0.count", None),
     }
   else:
     state_map = {
@@ -237,10 +238,18 @@ def get_layer_prefix(keystr_pax):
           f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
           transform_fn,
       )
+      state_map[f"['optimizer']['opt_state']['mu']{keystr_maxtext}.value"] = (
+          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
+          transform_fn,
+      )
       state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = (
           f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
           transform_fn,
       )
+      state_map[f"['optimizer']['opt_state']['nu']{keystr_maxtext}.value"] = (
+          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
+          transform_fn,
+      )
     else:
       state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn)
       state_map[f".opt_state.mu['params']{keystr_maxtext}"] = (