'tree_multimap" attribute error

[Guang]

I am getting this error when I try to run optimizers.py

module 'jax' has no attribute 'tree_multimap'

[Jannik Grothusen]

Fixed in #10.