FunMC: Use tree_util.tree_map instead of tree_util.tree_multimap to silence warning.
FunMC: Use tree_util.tree_map instead of tree_util.tree_multimap to silence warning.

PiperOrigin-RevId: 450797057

tensorflow/probability,tensorflow/probability