Conversion to unscanned format failed

### Bug report

hi, I first convert Llama HuggingFace checkpoints to Maxtext-compatible checkpoints using [llama_or_mistral_ckpt.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/MaxText/llama_or_mistral_ckpt.py) .
Then I use generate_param_only_checkpoint.py to convert the checkpoint into unscanned format. The error occurs here:
```
File "/project/maxtext/MaxText/generate_param_only_checkpoint.py", line 78, in unroll_layer_group
new_layer = jax.jit(slice_ith, out_shardings=new_layer_sharding)(layers)
File "/project/maxtext/MaxText/generate_param_only_checkpoint.py", line 94, in _possibly_unroll_params
unroll_layer_group(config.num_decoder_layers, layer_name="layers")
File "/project/maxtext/MaxText/generate_param_only_checkpoint.py", line 195, in generate_decode_checkpoint
_possibly_unroll_params(config, training_state, training_state_annotations, mesh)
File "/project/maxtext/MaxText/generate_param_only_checkpoint.py", line 211, in main
generate_decode_checkpoint(config)
File "/opt/miniconda3/envs/maxtext/lib/python3.10/site-packages/absl/app.py", line 261, in _run_main
sys.exit(main(argv))
File "/opt/miniconda3/envs/maxtext/lib/python3.10/site-packages/absl/app.py", line 316, in run
_run_main(main, args)
File "/Users/siqi/project/maxtext/MaxText/generate_param_only_checkpoint.py", line 215, in
app.run(main)
File "/opt/miniconda3/envs/maxtext/lib/python3.10/runpy.py", line 86, in _run_code
exec(code, run_globals)
File "/opt/miniconda3/envs/maxtext/lib/python3.10/runpy.py", line 196, in _run_module_as_main (Current frame)
return _run_code(code, main_globals, None,
TypeError: Argument 'input_layers['mlp']['mlp_layer_norm']['scale']' of shape float32[2048,16] of type is not a valid JAX type.
```

How can I fix this bug, or what's the best practise to run HuggingFace model?
thank you.

### Logs/Output

_No response_

### Environment Information

_No response_

### Additional Context

_No response_