A likely better way to mark parameters as non_trainable (at the array level rather than the tree level). At the tree level gives a simpler pytree definition, nicer printing etc, but the array level is more robust to avoiding missing attribute errors when e.g. manipulating pytrees.
A likely better way to mark parameters as non_trainable (at the array level rather than the tree level). At the tree level gives a simpler pytree definition, nicer printing etc, but the array level is more robust to avoiding missing attribute errors when e.g. manipulating pytrees.