orbax
orbax copied to clipboard
Add `tree_trim` function for filtering a PyTree by another's structure.
Add tree_trim function for filtering a PyTree by another's structure.