code: def make_variadic_reduce(reducer, vjp_bwd, tangents_fn):
      """Wraps a generic reducer function as a variadic reduction.
  
      The current use-case for this is TFP-internal. This function captures logic
      related to specific substrates and XLA, and enables some sharing of logic
      around `axis` and `keepdims` args.
  
      Args:
        reducer: The reducer callable. Takes two tuple args and returns a single
          reduced tuple.
        vjp_bwd: Custom VJP function. Takes `(aux, grads)` args with
          `aux = (operands, inits, axis, unsqueezed_shape)` and returns a tuple of
          grads w.r.t `operands`. Gradients w.r.t. `inits` are presumed `None`.
        tangents_fn: Custom JVP function. Takes `(inits, axis, primals, tangents)`
          args with `primals = (operands,)` and corresponding tangents and returns
          `tangents_out` (which must be linear w.r.t. `tangents`).
  
      Returns:
        reduce_fn: A callable with taking args
          `(operands, inits, axis=None, keepdims=False)`.
      """
  
      # Top-level `tf.function` for XLA (closed-over by the returned reduce_fn).
      @implementation_selection.never_runs_functions_eagerly
      @tf.function(jit_compile=True)
      def _xla_reduce(operands, inits, axis):
          """JIT-ed wrapper for TF `xla.variadic_reduce(..., reducer)`."""
          from tensorflow.compiler.tf2xla.python import (
              xla,
          )  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
  
          result = xla.variadic_reduce(
              operands,
              init_values=inits,
              dimensions_to_reduce=axis,
              reducer=tf.function(reducer).get_concrete_function(
                  inits, inits
              ),
          )
          # Graph mode: variadic reduce doesn't specify output shapes. Patch that.
          shp = operands[0].shape
          for arg in operands:
              shp = tensorshape_util.merge_with(shp, arg.shape)
          for part in result:
              tensorshape_util.set_shape(
                  part,
                  tuple(
                      dim for i, dim in enumerate(shp) if i not in axis
                  ),
              )
          return result
  
      def _variadic_reduce_no_grad(operands, inits, axis, reducer):
          if JAX_MODE:
              from jax import lax  # pylint: disable=g-import-not-at-top
  
              return lax.reduce(
                  operands,
                  init_values=inits,
                  dimensions=axis,
                  computation=reducer,
              )
          elif (
              tf.executing_eagerly()
              or not control_flow_util.GraphOrParentsInXlaContext(
                  tf1.get_default_graph()
              )
          ):
              return _variadic_reduce(
                  operands, init=inits, axis=axis, reducer=reducer
              )
          else:
              return _xla_reduce(operands, inits, axis)
  
      def _variadic_reduce_fwd(
          operands, inits, axis, reducer, unsqueezed_shape
      ):
          del unsqueezed_shape
          return (
              _variadic_reduce_no_grad(operands, inits, axis, reducer),
              (operands, inits),
          )
  
      def _variadic_reduce_jvp(
          axis, reducer, unsqueezed_shape, primals, tangents
      ):
          del unsqueezed_shape
          operands, inits = primals
          return (
              _variadic_reduce_no_grad(operands, inits, axis, reducer),
              tangents_fn(axis, primals, tangents),
          )
  
      @tfp_custom_gradient.custom_gradient(
          vjp_fwd=_variadic_reduce_fwd,
          vjp_bwd=vjp_bwd,
          jvp_fn=_variadic_reduce_jvp,
          nondiff_argnums=(2, 3, 4),
      )
      def _variadic_reduce_custom_grad(
          operands, inits, axis, reducer, unsqueezed_shape
      ):
          del unsqueezed_shape  # provided for backprop convenience
          return _variadic_reduce_no_grad(
              operands, inits, axis, reducer
          )
  
      def reduce_fn(operands, inits, axis=None, keepdims=False):
          """Applies `reducer` to the given operands along the given axes.
  
          Args:
            operands: tuple of tensors, all having the same shape.
            inits: tuple of scalar tensors, with dtypes aligned to those of operands.
            axis: The axis or axes to reduce. One of `None`, an `int` or a sequence of
              `int`. `None` is taken to mean "reduce all axes".
            keepdims: When `True`, we do not squeeze away the reduced dims, instead
              returning values with singleton dims in those axes.
  
          Returns:
            reduced: A tuple of the reduced operands.
          """
          # Static shape consistency checks.
          args_shape = operands[0].shape
          for arg in operands[1:]:
              args_shape = tensorshape_util.merge_with(
                  args_shape, arg.shape
              )
          ndims = tensorshape_util.rank(args_shape)
          if ndims is None:
              raise ValueError(
                  "Rank of at least one of `operands` must be known statically."
              )
          # Ensure the 'axis' arg is a tuple of non-negative ints.
          axis = np.arange(ndims) if axis is None else np.array(axis)
          if axis.ndim > 1:
              raise ValueError(
                  "`axis` must be `None`, an `int`, or a sequence of "
                  "`int`, but got {}".format(axis)
              )
          axis = np.reshape(axis, [-1])
          axis = np.where(axis < 0, axis + ndims, axis)
          axis = tuple(int(ax) for ax in axis)
  
          axis_nhot = ps.reduce_sum(
              ps.one_hot(
                  axis,
                  depth=ndims,
                  on_value=True,
                  off_value=False,
                  dtype=tf.bool,
              ),
              axis=0,
          )
          in_shape = args_shape
          if not tensorshape_util.is_fully_defined(in_shape):
              in_shape = tf.shape(operands[0])
          unsqueezed_shape = ps.where(axis_nhot, 1, in_shape)
  
          result = _variadic_reduce_custom_grad(
              operands, inits, axis, reducer, unsqueezed_shape
          )
  
          if keepdims:
              result = tf.nest.map_structure(
                  lambda t: tf.reshape(t, unsqueezed_shape), result
              )
          return result
  
      return reduce_fn
  
file_path: probability/tensorflow_probability/python/internal/variadic_reduce.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import implementation_selection
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 2 for calling a global function with example usage: make_variadic_reduce(reducer, vjp_bwd, tangents_fn) and returns: reduce_fn, result, _variadic_reduce_no_grad,  operands,  inits, _variadic_reduce_no_grad,  tangents_fn, _variadic_reduce_no_grad, result, lax, _variadic_reduce, _xla_reduce