flax

[nnx] register_variable_name refactor

#4540

Merged