code: def make_transforms(
      net: bnn.BNN,
  ) -> Tuple[
      Callable[..., Any], Callable[..., Any], Callable[..., Any]
  ]:
      """Returns unconstraining bijectors for all variables in the BNN."""
      jb = jax.tree_map(
          lambda x: x.experimental_default_event_space_bijector(),
          net.get_all_distributions(),
          is_leaf=lambda x: isinstance(
              x, distribution_lib.Distribution
          ),
      )
  
      def transform(params):
          return {
              "params": jax.tree_map(
                  lambda p, b: b(p), params["params"], jb
              )
          }
  
      def inverse_transform(params):
          return {
              "params": jax.tree_map(
                  lambda p, b: b.inverse(p), params["params"], jb
              )
          }
  
      def inverse_log_det_jacobian(params):
          return jax.tree_util.tree_reduce(
              lambda a, b: a + b,
              jax.tree_map(
                  lambda p, b: jnp.sum(b.inverse_log_det_jacobian(p)),
                  params["params"],
                  jb,
              ),
              initializer=0.0,
          )
  
      return transform, inverse_transform, inverse_log_det_jacobian
  
file_path: probability/tensorflow_probability/python/experimental/autobnn/util.py
imports: from typing import Any, Callable, List, Tuple
  import jax
  import jax.numpy
  import scipy
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.substrates.jax.distributions import distribution
label_desc: Use image node_id 1 for calling a global function with example usage: make_transforms(net) and returns: transform,  inverse_transform,  inverse_log_det_jacobian, dict, dict, jax