code: def make_psum_function(fn, in_axes, out_axes, out_dtype):
      """Constructs a function that broadcasts inputs over named axes.
  
      Given a function `fn`, `make_psum_function` returns a new one that
      includes psums over terms according to axis names provided in `out_axes`. It
      also adds psums for the vector-Jacobian product of the outputs of `fn` w.r.t.
      its inputs according to `in_axes` if there are axes in the outputs that are
      not present in an input.
  
      Args:
        fn: a callable to be transformed to have psums at its outputs and on the
          gradients to 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 outputs of `fn` will be psum-med according to their
          respective output axes.
        out_dtype: A structure of dtypes that matches the output of `fn`.
  
      Returns:
        A new function that applies psums on to the output of the original
        function and corrects the gradient with respect to its inputs.
      """
  
      out_axes = nest.map_structure_up_to(
          out_dtype, canonicalize_named_axis, out_axes
      )
  
      def psum_fn(*args):
          out = make_pbroadcast_function(
              fn, in_axes, out_axes, out_dtype
          )(*args)
  
          def _psum_output(x, out_axis):
              return psum(x, named_axis=out_axis)
  
          return nest.map_structure_up_to(
              out_dtype, _psum_output, out, out_axes
          )
  
      return psum_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 19 for calling a global function with example usage: make_psum_function(fn, in_axes, out_axes, out_dtype) and returns: psum_fn, nest, psum