code: def make_pbroadcast_function(fn, in_axes, out_axes, out_dtype):
      """Constructs a function that broadcasts inputs over named axes.
  
      Given a function `fn`, `make_pbroadcast_function` returns a new one that
      applies `pbroadcast` to input terms according to axis names provided in
      `in_axes` and `out_axes`. For each output axis in each term out the output of
      `fn`, inputs that do not have the output axes present are pbroadcasted before
      that term is computed.
  
      Args:
        fn: a callable to be transformed to have proadcasts at its inputs.
        in_axes: A structure of axis names that should match the structure of the
          input to `fn`. If the set of input axes for an input value does not match
          the output axes of a particular output value, the gradient of that output
          value w.r.t. the input value will be psum-ed over the axes present in the
          output but not the input.
        out_axes: A structure of axis names that should match the structure of the
          output of `fn`. The inputs to `fn` will be pbroadcast-ed before computing
          output terms according to their output axes.
        out_dtype: A structure of dtypes that matches the output of `fn`.
  
      Returns:
        A new function that applies pbroadcasts to the inputs of the original
        function.
      """
  
      if not isinstance(in_axes, tuple):
          in_axes = (in_axes,)
  
      def pbroadcast_fn(*args):
          nest.assert_shallow_structure(args, in_axes)
          nest.assert_shallow_structure(out_dtype, out_axes)
          map_in_axes = nest.map_structure_up_to(
              args, canonicalize_named_axis, in_axes
          )
          map_out_axes = nest.map_structure_up_to(
              out_dtype, canonicalize_named_axis, out_axes
          )
  
          def _pbroadcast_input(out_axes, x, in_axes):
              psum_axes = [
                  axis_name
                  for axis_name in out_axes
                  if axis_name not in in_axes
              ]
              return pbroadcast(x, psum_axes)
  
          def _flat_fn_index(i, *args):
              out = fn(*args)
              return tf.nest.flatten(out)[i]
  
          def _flat_fn(*args):
              outputs = []
              for i, out_axis in enumerate(
                  nest.flatten_up_to(out_dtype, map_out_axes)
              ):
                  local_args = nest.map_structure_up_to(
                      args,
                      functools.partial(_pbroadcast_input, out_axis),
                      args,
                      map_in_axes,
                  )
                  outputs.append(_flat_fn_index(i, *local_args))
              return tf.nest.pack_sequence_as(out_dtype, outputs)
  
          return _flat_fn(*args)
  
      return pbroadcast_fn
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow.python.util import nest
label_desc: Use image node_id 18 for calling a global function with example usage: make_pbroadcast_function(fn, in_axes, out_axes, out_dtype) and returns: pbroadcast_fn, _flat_fn, pbroadcast, tf, tf