diff --git a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py
index 2032cf1c41..098bd403f1 100644
--- a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py
+++ b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py
@@ -251,9 +251,12 @@ def tensor_getter(key):
   return {"params": restored_params}
 
 
-def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config):
+def load_safetensors_dynamic_state(path, abstract_params, maxtext_config):
   """Main entry point to dynamically build and load safetensors into MaxText format.
 
+  `abstract_params` is the weights of the target state -- Linen's `params` collection, or the
+  NNX params state -- not the full train state; the HF param mappings name weights only.
+
   Splits execution into:
   1. Deriving Mappings
   2. Loading Sharded arrays directly to TPUs
@@ -349,11 +352,7 @@ def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_con
   param_map_mt_to_hf, hook_fn_map_mt = get_hf_config_and_mappings(maxtext_config)
   max_logging.log(f"[1/3] Mappings derived in {time.time() - t_total:.2f}s")
 
-  target_tree = (
-      abstract_unboxed_pre_state.to_pure_dict()
-      if isinstance(abstract_unboxed_pre_state, nnx.State)
-      else abstract_unboxed_pre_state.params
-  )
+  target_tree = abstract_params.to_pure_dict() if isinstance(abstract_params, nnx.State) else abstract_params
 
   t1 = time.time()
   hf_state = load_sharded_hf_state(path)
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index cf51083597..7860386e72 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/src/maxtext/common/checkpointing.py
@@ -16,6 +16,7 @@
 """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer."""
 
 import datetime
+import importlib
 import time
 from typing import Any
 
@@ -157,33 +158,58 @@ def _load_linen_checkpoint_into_nnx(
   restore_args = ocp.checkpoint_utils.construct_restore_args(linen_abstract)
   restored = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True)
   restored = ckptr.restore(epath.Path(path), args=restored)
-  _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored))
-  return _linen_items_to_nnx(restored, abstract_nnx_state)
+  return _restored_linen_to_nnx(restored, abstract_nnx_state)
 
 
-def _restore_emergency_linen_checkpoint_into_nnx(
-    checkpoint_manager,
-    step,
-    abstract_nnx_state,
-    map_to_pspec,
-):
-  """Restores an emergency Linen-layout checkpoint into an NNX state.
+def _restored_linen_to_nnx(restored_linen, abstract_nnx_state):
+  """Reshapes a restored Linen-layout tree into the NNX state.
 
-  The `nnx_aux` subtree is stored inside `items`, so an emergency checkpoint
-  carries it too; it's restored when present and otherwise kept at its fresh
-  init value. A genuinely-missing weight raises.
+  Raises if the checkpoint is missing a weight. Every NNX restore path ends here: the load
+  itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout.
   """
-  max_logging.log(f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}")
-  linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state)
-  restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract)
-  checkpoint_args = ocp.args.PyTreeRestore(
-      item=linen_abstract,
-      restore_args=restore_args,
-      partial_restore=True,
-  )
-  restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state
-  _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored))
-  return _linen_items_to_nnx(restored, abstract_nnx_state)
+  _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored_linen))
+  return _linen_items_to_nnx(restored_linen, abstract_nnx_state)
+
+
+def _abstract_params(abstract_unboxed_pre_state):
+  """Returns the state's weights: the NNX Param subtree, or Linen's `params` collection."""
+  if isinstance(abstract_unboxed_pre_state, nnx.State):
+    return nnx.split_state(abstract_unboxed_pre_state.model, nnx.Param, ...)[0]
+  return abstract_unboxed_pre_state.params
+
+
+def _bare_weights(tree):
+  """Strips the Flax `params` collection wrapper so weights compare at the same depth.
+
+  A Linen params tree is the collection, an NNX one the bare weights; the dynamic
+  safetensors loader always returns the collection.
+  """
+  return tree["params"] if isinstance(tree, dict) and len(tree) == 1 and "params" in tree else tree
+
+
+def _resolve_conversion_fn(checkpoint_conversion_fn):
+  """Returns `checkpoint_conversion_fn` as a callable.
+
+  Config carries it as a dotted string ("my_pkg.my_module.my_fn"), so it has to be imported
+  before it can be called. A callable is used as is.
+  """
+  if checkpoint_conversion_fn is None:
+    raise ValueError(
+        "source_checkpoint_layout='safetensors' needs `checkpoint_conversion_fn` to map the "
+        "checkpoint's weights onto the model's, e.g. checkpoint_conversion_fn=my_pkg.my_module.my_fn."
+    )
+  if callable(checkpoint_conversion_fn):
+    return checkpoint_conversion_fn
+  module_name, _, fn_name = str(checkpoint_conversion_fn).rpartition(".")
+  if not module_name:
+    raise ValueError(f"`checkpoint_conversion_fn` must be a dotted path to a function, got {checkpoint_conversion_fn!r}.")
+  try:
+    fn = getattr(importlib.import_module(module_name), fn_name, None)
+  except ImportError as e:
+    raise ValueError(f"Could not import `checkpoint_conversion_fn` {checkpoint_conversion_fn!r}: {e}") from e
+  if not callable(fn):
+    raise ValueError(f"`checkpoint_conversion_fn` {checkpoint_conversion_fn!r} is not a function.")
+  return fn
 
 
 def _load_full_state_from_path(
@@ -226,6 +252,8 @@ def _load_full_state_from_path(
       with context:
         return ocp_v1.load_pytree(path, abstract_unboxed_pre_state)
     elif source_checkpoint_layout == "safetensors":
+      # Resolved first, so a bad config fails before the weights are read.
+      conversion_fn = _resolve_conversion_fn(checkpoint_conversion_fn)
       context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.SAFETENSORS)
       with context:
         metadata = ocp_v1.pytree_metadata(path)
@@ -237,7 +265,11 @@ def combine_sharding(sds, shardings):
 
         sharded_abstract_state = jax.tree.map(combine_sharding, simple_abstract_state, shardings)
         pre_transformed_state = ocp_v1.load_pytree(path, sharded_abstract_state)
-      state = checkpoint_conversion_fn(pre_transformed_state)
+      state = conversion_fn(pre_transformed_state)
+      # The conversion fn returns MaxText's on-disk (Linen) layout, which is what pure_nnx reads,
+      # so NNX needs the same reshape as every other restore. An NNX state passes through.
+      if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(state, nnx.State):
+        state = _restored_linen_to_nnx(state, abstract_unboxed_pre_state)
       return state
     else:
       raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}")
@@ -406,6 +438,11 @@ def load_state_if_possible(
      set.
   """
 
+  # pure_nnx saves in the Linen on-disk layout, so every branch below loads the same tree Linen
+  # does: the NNX abstract is converted to that layout going in, and what comes back is reshaped
+  # into the NNX state on the way out.
+  is_nnx = isinstance(abstract_unboxed_pre_state, nnx.State)
+
   if checkpoint_manager is not None:
     max_logging.log("checkpoint manager exists so trying to load this run's existing checkpoint")
 
@@ -437,49 +474,9 @@ def map_to_pspec(data):
         )
         ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True)
 
-      # pure_nnx saves in the Linen on-disk layout; restore that layout (weights +
-      # opt_state + step + nnx_aux), restoring the grain iterator in place when
-      # present, then reshape it back into the NNX state.
-      # (Emergency managers use their own restore path below.)
-      if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(
-          checkpoint_manager,
-          (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager),
-      ):
-        linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state)
-        restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract)
-        checkpoint_args = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True)
-        if (
-            dataset_type == "grain"
-            and data_iterator
-            and not isinstance(data_iterator, PlaceHolderDataIterator)
-            and (checkpoint_manager.directory / str(step) / "iter").exists()
-        ):
-          restored, _ = grain_utility.restore_grain_iterator(
-              checkpoint_manager, step, data_iterator, checkpoint_args, expansion_factor_real_data
-          )
-        else:
-          restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args))
-        _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_unboxed_pre_state, restored["items"]))
-        restored_nnx = _linen_items_to_nnx(restored["items"], abstract_unboxed_pre_state)
-        return ({"items": restored_nnx}, None)
-
-      if isinstance(abstract_unboxed_pre_state, nnx.State) and isinstance(
-          checkpoint_manager,
-          (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager),
-      ):
-        restored = _restore_emergency_linen_checkpoint_into_nnx(
-            checkpoint_manager,
-            step,
-            abstract_unboxed_pre_state,
-            map_to_pspec,
-        )
-        return (
-            restored,
-            None,
-        )
-
-      # Only Linen TrainState reaches here; the NNX cases returned above.
-      restore_target = abstract_unboxed_pre_state
+      restore_target = (
+          train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state) if is_nnx else abstract_unboxed_pre_state
+      )
       restore_args = jax.tree_util.tree_map(map_to_pspec, restore_target)
       checkpoint_args = ocp.args.PyTreeRestore(
           item=restore_target,
@@ -499,6 +496,8 @@ def map_to_pspec(data):
             ),
         ):
           restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state
+          if is_nnx:
+            restored = _restored_linen_to_nnx(restored, abstract_unboxed_pre_state)
           return (
               restored,
               None,
@@ -515,30 +514,43 @@ def map_to_pspec(data):
             and not isinstance(data_iterator, PlaceHolderDataIterator)
             and (checkpoint_manager.directory / str(step) / "iter").exists()
         ):
-          return grain_utility.restore_grain_iterator(
+          restored, iterator = grain_utility.restore_grain_iterator(
               checkpoint_manager,
               step,
               data_iterator,
               checkpoint_args,
               expansion_factor_real_data,
           )
+          if is_nnx:
+            restored = {"items": _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state)}
+          return (restored, iterator)
         # Case 3: Default/Fallback case.
         # This case acts as a wildcard ('_') and matches if none of the preceding cases were met.
         case _:
           restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args))
+          if is_nnx:
+            restored = {"items": _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state)}
           return (restored, None)
 
   if source_checkpoint_layout == "safetensors_dynamic":
     path = load_parameters_from_path or load_full_state_from_path
     max_logging.log(f"Dynamic On-the-Fly Formatting: Loading SafeTensors from {path}")
 
-    return load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config)
+    # Weights-only for both paths, so the loader gets the weights rather than the whole state:
+    # the HF param mappings name weights, and an NNX state hides them under `model`.
+    params = _abstract_params(abstract_unboxed_pre_state)
+    restored, restored_params = load_safetensors_dynamic_state(path, params, maxtext_config)
+    # A weight no HF mapping covered comes back unmaterialized and would reach the model as an
+    # untrained init value. Same check the Orbax weights-only load makes.
+    _raise_on_weight_mismatch(_bare_weights(params.to_pure_dict() if is_nnx else params), _bare_weights(restored_params))
+    if is_nnx:
+      # The loader returns the Linen `params` collection; NNX holds bare weights, so unwrap it
+      # back into the params state, the shape load_params_from_path returns.
+      nnx.replace_by_pure_dict(params, restored_params["params"])
+      return restored, params
+    return restored, restored_params
   elif load_parameters_from_path != "":
-    if isinstance(abstract_unboxed_pre_state, nnx.State):
-      _, params, _ = nnx.split(abstract_unboxed_pre_state.model, nnx.Param, ...)
-    else:
-      params = abstract_unboxed_pre_state.params
-
+    params = _abstract_params(abstract_unboxed_pre_state)
     restored_params = load_params_from_path(
         load_parameters_from_path,
         params,
diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py
index 163a5b1e72..0c2dae51ff 100644
--- a/tests/unit/checkpointing_nnx_load_test.py
+++ b/tests/unit/checkpointing_nnx_load_test.py
@@ -21,6 +21,7 @@
 
 from etils import epath
 from flax import nnx
+from flax.training import train_state
 import jax
 import jax.numpy as jnp
 from maxtext.common import checkpointing
@@ -83,13 +84,16 @@ def test_emergency_restore_recovers_nnx_aux(self):
     saved_linen = train_state_nnx.to_checkpoint_dict(concrete)
     self.assertIn("nnx_aux", saved_linen)
 
-    checkpoint_manager = mock.Mock()
+    checkpoint_manager = mock.Mock(spec=checkpointing.EmergencyCheckpointManager)
     checkpoint_manager.restore.return_value = mock.Mock(state=saved_linen)
-    restored = checkpointing._restore_emergency_linen_checkpoint_into_nnx(  # pylint: disable=protected-access
-        checkpoint_manager,
-        14,
-        _abstract_dropout_state(),
-        lambda leaf: ocp.type_handlers.ArrayRestoreArgs(global_shape=leaf.shape, dtype=leaf.dtype),
+    restored, _ = checkpointing.load_state_if_possible(
+        checkpoint_manager=checkpoint_manager,
+        data_iterator=None,
+        load_parameters_from_path="",
+        load_full_state_from_path="",
+        checkpoint_storage_concurrent_gb=8,
+        abstract_unboxed_pre_state=_abstract_dropout_state(),
+        step=14,
     )
     restored = restored.to_pure_dict()
 
@@ -147,6 +151,74 @@ def test_load_parameters_from_path_uses_state_params_for_linen(self):
     forwarded_params = m.call_args[0][1]
     self.assertIs(forwarded_params, fake_state.params)
 
+  def test_safetensors_dynamic_hands_over_weights_and_returns_nnx_params(self):
+    """The loader gets the bare weights, and its result comes back as NNX params."""
+    abstract = _abstract_nnx_state()
+    weights = {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.array([5.0])}}
+
+    with mock.patch.object(
+        checkpointing, "load_safetensors_dynamic_state", return_value=(None, {"params": weights})
+    ) as m:
+      full, params = checkpointing.load_state_if_possible(
+          checkpoint_manager=None,
+          data_iterator=None,
+          load_parameters_from_path="gs://does-not-exist/hf",
+          load_full_state_from_path="",
+          checkpoint_storage_concurrent_gb=8,
+          abstract_unboxed_pre_state=abstract,
+          source_checkpoint_layout="safetensors_dynamic",
+          maxtext_config=mock.Mock(),
+      )
+
+    self.assertIsNone(full)
+    # The loader sees the weights the way the HF mappings name them: no `model` prefix above them.
+    forwarded = m.call_args[0][1]
+    self.assertEqual(set(forwarded.to_pure_dict().keys()), {"linear"})
+    # And the restored weights come back as an nnx.State, the shape a params-only load returns.
+    self.assertIsInstance(params, nnx.State)
+    self.assertTrue(jnp.array_equal(params.to_pure_dict()["linear"]["bias"], weights["linear"]["bias"]))
+
+  def test_safetensors_dynamic_unmapped_weight_raises(self):
+    """A weight no HF mapping covered comes back unmaterialized, which must raise."""
+    unmapped = {
+        "linear": {"kernel": jnp.ones((2, 1)), "bias": jax.ShapeDtypeStruct((1,), jnp.float32)},
+    }
+    with mock.patch.object(checkpointing, "load_safetensors_dynamic_state", return_value=(None, {"params": unmapped})):
+      with self.assertRaises(ValueError) as ctx:
+        checkpointing.load_state_if_possible(
+            checkpoint_manager=None,
+            data_iterator=None,
+            load_parameters_from_path="gs://does-not-exist/hf",
+            load_full_state_from_path="",
+            checkpoint_storage_concurrent_gb=8,
+            abstract_unboxed_pre_state=_abstract_nnx_state(),
+            source_checkpoint_layout="safetensors_dynamic",
+            maxtext_config=mock.Mock(),
+        )
+    self.assertIn("linear/bias", str(ctx.exception))
+
+  def test_safetensors_dynamic_forwards_params_collection_for_linen(self):
+    """Linen hands over its `params` collection, not the whole TrainState, and gets it back as-is."""
+    weights = {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}
+    fake_state = mock.Mock(spec=["params"])
+    fake_state.params = weights
+
+    with mock.patch.object(checkpointing, "load_safetensors_dynamic_state", return_value=(None, weights)) as m:
+      full, params = checkpointing.load_state_if_possible(
+          checkpoint_manager=None,
+          data_iterator=None,
+          load_parameters_from_path="gs://does-not-exist/hf",
+          load_full_state_from_path="",
+          checkpoint_storage_concurrent_gb=8,
+          abstract_unboxed_pre_state=fake_state,
+          source_checkpoint_layout="safetensors_dynamic",
+          maxtext_config=mock.Mock(),
+      )
+
+    self.assertIsNone(full)
+    self.assertIs(m.call_args[0][1], fake_state.params)
+    self.assertIs(params, weights)  # Linen keeps the collection layout it passed in
+
   def test_no_paths_returns_none_none(self):
     """Sanity: with no checkpoint manager and no load paths, the function returns (None, None)."""
     full, params = checkpointing.load_state_if_possible(
@@ -161,6 +233,158 @@ def test_no_paths_returns_none_none(self):
     self.assertIsNone(params)
 
 
+class TestManagerRestoreParity(unittest.TestCase):
+  """NNX and Linen go through the same manager restore; only the ends of the call differ.
+
+  Cases come in pairs, one per state type, so a divergence in the shared flow fails one half.
+  """
+
+  def _manager(self, restored):
+    manager = mock.MagicMock(spec=ocp.CheckpointManager)
+    manager.restore.return_value = restored
+    return manager
+
+  def _linen_abstract(self):
+    """A real Linen TrainState -- the manager restore paths tree_map over the whole state."""
+
+    def make():
+      params = {"params": {"linear": {"kernel": jnp.zeros((2, 1))}}}
+      return train_state.TrainState.create(apply_fn=lambda x: x, params=params, tx=optax.identity())
+
+    return jax.eval_shape(make)  # every leaf abstract, `step` included, as get_abstract_state gives it
+
+  def _load(self, manager, abstract, **kwargs):
+    kwargs.setdefault("data_iterator", None)
+    return checkpointing.load_state_if_possible(
+        checkpoint_manager=manager,
+        load_parameters_from_path="",
+        load_full_state_from_path="",
+        checkpoint_storage_concurrent_gb=8,
+        abstract_unboxed_pre_state=abstract,
+        step=7,
+        **kwargs,
+    )
+
+  def test_nnx_restores_the_linen_layout_and_returns_an_nnx_state(self):
+    abstract = _abstract_dropout_state()
+    saved = train_state_nnx.to_checkpoint_dict(_dropout_state(0))
+    manager = self._manager({"items": saved})
+
+    restored, _ = self._load(manager, abstract)
+
+    # Going in: the manager is asked for the Linen on-disk layout, not the NNX one.
+    item = manager.restore.call_args.kwargs["args"]["items"].item
+    self.assertEqual(set(item) & {"params", "step"}, {"params", "step"})
+    self.assertNotIn("model", item)
+    # Coming out: reshaped back under `items`, the same key the Linen path returns.
+    self.assertIsInstance(restored["items"], nnx.State)
+    self.assertIn("model", restored["items"].to_pure_dict())
+
+  def test_linen_restore_target_and_return_are_untouched(self):
+    abstract = self._linen_abstract()
+    sentinel = {"items": {"params": abstract.params}}
+    manager = self._manager(sentinel)
+
+    restored, _ = self._load(manager, abstract)
+
+    self.assertIs(manager.restore.call_args.kwargs["args"]["items"].item, abstract)
+    self.assertIs(restored, sentinel)  # returned exactly as the manager gave it
+
+  def test_grain_case_converts_items_and_passes_the_iterator_through(self):
+    """The grain branch is shared too: NNX only reshapes `items`, leaving the iterator element alone."""
+    abstract = _abstract_nnx_state()
+    saved = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.array([5.0])}}}}
+    manager = self._manager(None)
+
+    with mock.patch.object(
+        checkpointing.grain_utility, "restore_grain_iterator", return_value=({"items": saved}, None)
+    ) as m:
+      restored, iterator = self._load(manager, abstract, dataset_type="grain", data_iterator=mock.MagicMock())
+
+    m.assert_called_once()
+    self.assertIsNone(iterator)
+    self.assertIsInstance(restored["items"], nnx.State)
+    self.assertTrue(jnp.array_equal(restored["items"].to_pure_dict()["model"]["linear"]["bias"], jnp.array([5.0])))
+
+  def test_grain_case_is_untouched_for_linen(self):
+    abstract = self._linen_abstract()
+    sentinel = ({"items": {"params": abstract.params}}, None)
+    manager = self._manager(None)
+
+    with mock.patch.object(checkpointing.grain_utility, "restore_grain_iterator", return_value=sentinel):
+      restored, iterator = self._load(manager, abstract, dataset_type="grain", data_iterator=mock.MagicMock())
+
+    self.assertIs(restored, sentinel[0])
+    self.assertIsNone(iterator)
+
+  def test_missing_weight_raises_on_the_standard_path(self):
+    """The missing-weight check guards every NNX restore, not just the weights-only load."""
+    saved = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}}  # no bias
+    manager = self._manager({"items": saved})
+
+    with self.assertRaises(ValueError) as ctx:
+      self._load(manager, _abstract_nnx_state())
+    self.assertIn("linear/bias", str(ctx.exception))
+
+  def test_missing_weight_raises_on_the_emergency_path(self):
+    saved = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}}  # no bias
+    manager = mock.Mock(spec=checkpointing.EmergencyCheckpointManager)
+    manager.restore.return_value = mock.Mock(state=saved)
+
+    with self.assertRaises(ValueError) as ctx:
+      self._load(manager, _abstract_nnx_state())
+    self.assertIn("linear/bias", str(ctx.exception))
+
+  def test_no_step_in_manager_falls_through_to_the_load_paths(self):
+    """An empty manager (latest_step() is None) must not restore -- it falls through, for both types."""
+    manager = mock.MagicMock(spec=ocp.CheckpointManager)
+    manager.latest_step.return_value = None
+
+    for abstract in (_abstract_nnx_state(), self._linen_abstract()):
+      restored, params = checkpointing.load_state_if_possible(
+          checkpoint_manager=manager,
+          data_iterator=None,
+          load_parameters_from_path="",
+          load_full_state_from_path="",
+          checkpoint_storage_concurrent_gb=8,
+          abstract_unboxed_pre_state=abstract,
+      )
+      self.assertIsNone(restored)
+      self.assertIsNone(params)
+      manager.restore.assert_not_called()
+
+
+class TestResolveConversionFn(unittest.TestCase):
+  """`checkpoint_conversion_fn` arrives from config as a dotted string and has to be imported."""
+
+  def test_dotted_string_is_imported(self):
+    fn = checkpointing._resolve_conversion_fn("json.loads")  # pylint: disable=protected-access
+    self.assertEqual(fn('{"a": 1}'), {"a": 1})
+
+  def test_callable_passes_through(self):
+    self.assertIs(checkpointing._resolve_conversion_fn(len), len)  # pylint: disable=protected-access
+
+  def test_missing_fn_names_the_config_flag(self):
+    with self.assertRaises(ValueError) as ctx:
+      checkpointing._resolve_conversion_fn(None)  # pylint: disable=protected-access
+    self.assertIn("checkpoint_conversion_fn", str(ctx.exception))
+
+  def test_bare_name_without_module_raises(self):
+    with self.assertRaises(ValueError) as ctx:
+      checkpointing._resolve_conversion_fn("my_fn")  # pylint: disable=protected-access
+    self.assertIn("dotted path", str(ctx.exception))
+
+  def test_unimportable_module_raises(self):
+    with self.assertRaises(ValueError) as ctx:
+      checkpointing._resolve_conversion_fn("no_such_module_xyz.my_fn")  # pylint: disable=protected-access
+    self.assertIn("no_such_module_xyz", str(ctx.exception))
+
+  def test_non_callable_attribute_raises(self):
+    with self.assertRaises(ValueError) as ctx:
+      checkpointing._resolve_conversion_fn("json.__doc__")  # pylint: disable=protected-access
+    self.assertIn("not a function", str(ctx.exception))
+
+
 class TestLoadParamsIntoNNX(unittest.TestCase):
   """Weight-only load (load_parameters_path) of a Linen-layout checkpoint into NNX."""
 
@@ -214,6 +438,52 @@ def test_linen_params_dict_restores_as_the_params_collection(self):
     self.assertTrue(jnp.array_equal(restored["params"]["linear"]["bias"], weights["linear"]["bias"]))
 
 
+class TestSafetensorsFullStateIntoNNX(unittest.TestCase):
+  """A full-state safetensors load reshapes the conversion fn's output into the NNX state."""
+
+  def _load(self, converted, abstract):
+    """Runs the v1 safetensors branch, stubbing the read so only the conversion is under test."""
+    with (
+        mock.patch.object(checkpointing, "ocp_v1") as v1,
+        mock.patch.object(checkpointing, "sharding_utils") as shardings,
+    ):
+      v1.pytree_metadata.return_value = mock.Mock(metadata={"w": jax.ShapeDtypeStruct((1,), jnp.float32)})
+      shardings.construct_maximal_shardings.return_value = {"w": None}
+      return checkpointing._load_full_state_from_path(  # pylint: disable=protected-access
+          path="gs://does-not-exist/hf",
+          abstract_unboxed_pre_state=abstract,
+          enable_orbax_v1=True,
+          checkpoint_conversion_fn=lambda _: converted,
+          source_checkpoint_layout="safetensors",
+          checkpoint_storage_concurrent_gb=8,
+          use_ocdbt=True,
+          use_zarr3=True,
+      )
+
+  def test_linen_layout_output_is_reshaped_into_nnx(self):
+    converted = {
+        "params": {"params": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.array([5.0])}}},
+        "step": jnp.asarray(0, jnp.int32),
+    }
+    restored = self._load(converted, _abstract_nnx_state())
+
+    self.assertIsInstance(restored, nnx.State)
+    pure = restored.to_pure_dict()
+    self.assertTrue(jnp.array_equal(pure["model"]["linear"]["kernel"], jnp.ones((2, 1))))
+
+  def test_missing_weight_raises(self):
+    """A bad conversion fn hits the same missing-weight check as the other NNX loads."""
+    converted = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}}  # no bias
+    with self.assertRaises(ValueError) as ctx:
+      self._load(converted, _abstract_nnx_state())
+    self.assertIn("linear/bias", str(ctx.exception))
+
+  def test_linen_state_is_untouched(self):
+    """For Linen the conversion fn's output is returned as-is."""
+    converted = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}}
+    self.assertIs(self._load(converted, mock.Mock(spec=["params"])), converted)
+
+
 class TestToCheckpointDict(unittest.TestCase):
   """train_state_nnx.to_checkpoint_dict splits the NNX state by Variable type."""
 
diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py
index f21f025175..421a6d150f 100644
--- a/tests/unit/checkpointing_test.py
+++ b/tests/unit/checkpointing_test.py
@@ -228,15 +228,8 @@ def __init__(self):
 
     config = MockConfig()
 
-    class DummyAbstractState:
-
-      def __init__(self):
-        self.params = {}
-
-    abstract_state = DummyAbstractState()
-
     path = "repo/meta-llama"
-    dummy_ret_val, loaded_vars = load_dynamic.load_safetensors_dynamic_state(path, abstract_state, config)
+    dummy_ret_val, loaded_vars = load_dynamic.load_safetensors_dynamic_state(path, {}, config)
 
     self.assertIsNone(dummy_ret_val)
     self.assertEqual(loaded_vars, {"params": {}})
diff --git a/tests/unit/setup_initial_state_nnx_test.py b/tests/unit/setup_initial_state_nnx_test.py
index 0f6477c9bb..8130f349fd 100644
--- a/tests/unit/setup_initial_state_nnx_test.py
+++ b/tests/unit/setup_initial_state_nnx_test.py
@@ -216,7 +216,7 @@ def test_emergency_manager_overlay_is_not_unwrapped(self):
     trained = _init_fn(_Model, 0)()
     grads = nnx.grad(lambda m: jnp.mean(m(_TRAIN_X, deterministic=False) ** 2))(trained.model)
     trained.apply_gradients(grads)
-    overlay = nnx.state(trained)  # what _restore_emergency_linen_checkpoint_into_nnx returns
+    overlay = nnx.state(trained)  # what an emergency restore returns: the NNX state, unwrapped
     manager = mock.Mock(spec=emergency_checkpoint_manager.CheckpointManager)
 
     with mock.patch.object(checkpointing, "load_state_if_possible", return_value=(overlay, None)):