flax

[nnx] try casting integers keys in State.replace_by_pure_dict

#4317

Merged