Replace jax.tree_util.tree_map() with jax.tree_util.tree_multimap()#3

oikosohn wants to merge 1 commit intogordicaleksa:maingordicaleksa/get-started-with-JAX:mainfrom

oikosohn:oikosohn-patch-1oikosohn/get-started-with-JAX:oikosohn-patch-1Copy head branch name to clipboard