code: def tree_map_distribution(
      f: Callable[
          [mpo_types.DistributionOrArray], mpo_types.DistributionOrArray
      ],
      x: mpo_types.DistributionOrArray,
  ) -> mpo_types.DistributionOrArray:
      """Apply a jax function to a distribution by treating it as tree."""
      if isinstance(x, distrax.Distribution):
          safe_f = lambda y: f(y) if isinstance(y, jnp.ndarray) else y
          nil, tree_data = x.tree_flatten()
          new_tree_data = jax.tree_map(safe_f, tree_data)
          new_x = x.tree_unflatten(new_tree_data, nil)
          return new_x
      elif isinstance(x, tfd.Distribution):
          return jax.tree_map(f, x)
      else:
          return f(x)
  
file_path: acme/acme/agents/jax/mpo/utils.py
imports: from typing import Callable
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import types
  import distrax
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 4 for calling a global function with example usage: tree_map_distribution(f, x) and returns: new_x, jax, f