"LayerNorm", "scale"))
...     for path in flat_params
... }
>>> mask = traverse_util.unflatten_dict(mask)
>>> params = model.to_bf16(params, mask)  to_fp16  < source > ( params: Union mask: Any = None  )   Parameters   params (Union[Dict, FrozenDict]) —
A PyTree of model parameters.   mask (Union[Dict, FrozenDict]) —
A PyTree with same structure as the params tree. The leaves should be booleans. It should be True
for params you want to cast, a