updated occurences of tree_multimap to tree_map

[shyams2]

I was playing around with the colabs but noticed they were not working on account of tree_multimap being deprecated in the latest version of JAX. I've just updated every occurrence of tree_multimap with tree_map. With this change, the colab seems to work as intended.

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.