Loads and validates the weights. Returns a loaded subset of the weights.
(loader: _weight_loaders.WeightLoader, params_shape: at.Params)
| 71 | |
| 72 | |
| 73 | def _load_weights_and_validate(loader: _weight_loaders.WeightLoader, params_shape: at.Params) -> at.Params: |
| 74 | """Loads and validates the weights. Returns a loaded subset of the weights.""" |
| 75 | loaded_params = loader.load(params_shape) |
| 76 | at.check_pytree_equality(expected=params_shape, got=loaded_params, check_shapes=True, check_dtypes=True) |
| 77 | |
| 78 | # Remove jax.ShapeDtypeStruct from the loaded params. This makes sure that only the loaded params are returned. |
| 79 | return traverse_util.unflatten_dict( |
| 80 | {k: v for k, v in traverse_util.flatten_dict(loaded_params).items() if not isinstance(v, jax.ShapeDtypeStruct)} |
| 81 | ) |
| 82 | |
| 83 | |
| 84 | @at.typecheck |
no test coverage detected