flax

c67eb094 - [nnx] register_variable_name refactor

[nnx] register_variable_name refactor Renames `register_variable_name_type_pair` to `register_variable_name` and adds and overload that can be used as a class decorator to improve usability. PiperOrigin-RevId: 725592906