Flax模型用Safetensors保存？先把嵌套字典展平

Safetensors 文档里找不到“JAX”，但搜一下会指向一个 2023 年就没再更新的 GitHub 仓库。再仔细看，文档里其实有一个 Flax API 链接——名字起得有点歪，它本质上就是 JAX API。

这个 API 的 save_file 函数只接受 Dict[str, Array]——键是字符串，值是 JAX 数组。但 Flax 的 nnx.State.to_pure_dict 返回的是嵌套字典，不是单层字典。如果你直接把嵌套字典丢给 Safetensors，会报错：

` SafetensorError: dtype object is not covered `

原因很简单：Safetensors 没检查字典值的类型，以为每个值都是 JAX 数组，于是用 NumPy 的 asarray 去包装——asarray 能把任意对象包成 dtype=object 的数组，但底层的 Rust 代码不认识 object 类型，直接崩了。

HuggingFace 没有计划扩展 Safetensors 支持 FrozenDict 这类结构（见 #138）。所以社区有人写了 safejax 来处理序列化，但你也可以自己动手。

Flax 提供了 nnx.to_flat_state，它把模型状态变成 FlatState——一个列表，每个元素是 (键元组, Param) 对。Param 对象实现了 asarray 能识别的接口，直接转 NumPy 数组没问题。遍历这个列表，把键元组用点连接成字符串，就拼出了 Safetensors 需要的单层字典。

加载时反过来：把点分隔的键拆回元组，注意处理整数索引（比如 nnx.Sequential 里的索引）。最后用 nnx.from_flat_state 和 nnx.update 把字典恢复成模型状态。

三步：展平、转换、还原。不需要额外库。