fix module 'jax' has no attribute 'tree_multimap'#72

FlorenceZh wants to merge 2 commits into

magenta:mainmagenta/mt3:mainfrom

FlorenceZh:mainFlorenceZh/mt3:mainCopy head branch name to clipboard