Use jax.tree_util.tree_map in place of deprecated tree_multimap.#208

copybara-service[bot] wants to merge 1 commit intomaingoogle/learned_optimization:mainfrom

test_462218807google/learned_optimization:test_462218807Copy head branch name to clipboard