Use jax.tree_util.treedef_is_leaf for readability - #128

Use jax.tree_util.treedef_is_leaf for readability#128

findmyway wants to merge 1 commit into

google-deepmind:maingoogle-deepmind/penzai:mainfrom

findmyway:patch-1findmyway/penzai:patch-1Copy head branch name to clipboard