Internal

tests/models/safetensors_loader_test.py

@@ -82,8 +82,9 @@ def setUpClass(cls):
       cls.tensors[f'layers.{i}.w2.kernel'] = np.array(
           layer_state['w2']['kernel'].value
       )
-      cls.tensors[f'layers.{i}.w2.bias'] = np.array(
-          layer_state['w2']['bias'].value
+      # Test that nnx.Param are correctly handled.
+      cls.tensors[f'layers.{i}.w2.bias'] = nnx.Param(
+          np.array(layer_state['w2']['bias'].value),
       )
 
   @parameterized.named_parameters(

tunix/models/safetensors_loader.py

@@ -258,17 +258,23 @@ def process_key(k_name, f, sf_file, file_loaded_tensors):
     def make_update_tensor_fn(current_file_tensors):
       def update_tensor(path, param, shard=None):
         current_path_key = path_to_key(path)
-        if current_path_key in current_file_tensors:
-          loaded_arr = current_file_tensors[current_path_key]
-          if loaded_arr.shape != param.shape:
-            raise ValueError(
-                f'Shape mismatch for {current_path_key}: got'
-                f' {loaded_arr.shape}, expected {param.shape}'
-            )
-          if shard is not None:
-            return jax.device_put(loaded_arr, shard)
-          else:
-            return jax.device_put(loaded_arr, jax.devices()[0])
+
+        # nnx.Param adds a .value suffix to the key
+        possible_keys = [current_path_key, f'{current_path_key}.value']
+
+        for k in possible_keys:
+          if k in current_file_tensors:
+            loaded_arr = current_file_tensors[k]
+            if loaded_arr.shape != param.shape:
+              raise ValueError(
+                  f'Shape mismatch for {k}: got'
+                  f' {loaded_arr.shape}, expected {param.shape}'
+              )
+            if shard is not None:
+              return jax.device_put(loaded_arr, shard)
+            else:
+              return jax.device_put(loaded_arr, jax.devices()[0])
+
         return param
 
       return update_tensor