Use jax.tree_util.treedef_is_leaf for readability#128

findmyway wants to merge 1 commit intogoogle-deepmind:maingoogle-deepmind/penzai:mainfrom

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