Refactor

[cgarciae]

Changes

• TreeParts are now generic and statically behave like Union.
• filter now also accepts predicates.
• Optimizer.update was renamed to apply_updates.
• Expanded TreePart hierarchy.
• Added RngSeq Module (generates PRNGKeys on demand).
• TreeObject now has a metaclass that checks that super().__init__() is always called.
• TreeObjects are now automatically annotated if an annotation is not provided by the user.
• filter now also accepts predicates of the type FieldInfo -> bool.
• Added the Static annotation for when you want a field to be explicitly marked as a static part of the Pytree. This is useful if the field will hold a TreeObject but you don't want it to be a child of the Pytree e.g. ignored_linear: tx.Static[tx.Linear]