flax

[nnx] improve graph_utils._set_key_tuple

#3592

Merged