JAX backend doesn't use jax.tree_util

[krzysztofrusek]

In this part, expand_composites is set to False and this prevents jax.tree_util for kicking in.

probability/tensorflow_probability/python/internal/backend/numpy/nest.py

Lines 314 to 320 in 64a70d0

	def map_structure(func, *structure, **kwargs):
	"""Add expand_composites support for JAX."""
	expand_composites = kwargs.pop('expand_composites', False)
	if expand_composites and JAX_MODE:
	from jax import tree_util # pylint: disable=g-import-not-at-top
	return tree_util.tree_map(func, *structure)
	return dm_tree.map_structure(func, *structure, **kwargs)

The resulting problem is that we cannot use pytree incompatible with dm_tree

Herę is notebook reproducing the bug
https://colab.research.google.com/gist/krzysztofrusek/a9fa71ca2bf3952a9f18358309225107/eqx_tfp.ipynb

[SiegeLordEx]

expand_composites should probably be set to True in that location. There are other locations of similar nature where it is set to True (e.g.

probability/tensorflow_probability/python/internal/loop_util.py

Line 183 in 23a292a

initial_state, expand_composites=True)

,

probability/tensorflow_probability/python/mcmc/internal/util.py

Line 116 in 23a292a

flat_p = tf.nest.flatten(proposed, expand_composites=True)

), so I don't see an obvious reason why it shouldn't be True here as well.