code: def scan_associative(
      fn,
      elems,
      max_num_levels=48,
      axis=0,
      validate_args=False,
      name=None,
  ):
      """Perform a scan with an associative binary operation, in parallel.
  
      The associative scan operation computes the cumulative sum, or
      [all-prefix sum](https://en.wikipedia.org/wiki/Prefix_sum), of a set of
      elements under an associative binary operation [1]. For example, using the
      ordinary addition operator `fn = lambda a, b: a + b`, this is equivalent to
      the ordinary cumulative sum `tf.math.cumsum` along axis 0. This method
      supports the general case of arbitrary associative binary operations operating
      on `Tensor`s or structures of `Tensor`s:
  
      ```python
      scan_associative(fn, elems) = tf.stack([
        elems[0],
        fn(elems[0], elems[1]),
        fn(elems[0], fn(elems[1], elems[2])),
        ...
        fn(elems[0], fn(elems[1], fn(..., fn(elems[-2], elems[-1]))),
      ], axis=0)
      ```
  
      The associative structure allows the computation to be decomposed
      and executed by parallel reduction. Where a naive sequential
      implementation would loop over all `N` elements, this method requires
      only a logarithmic number (`2 * ceil(log_2 N)`) of sequential steps, and
      can thus yield substantial performance speedups from hardware-accelerated
      vectorization. The total number of invocations of the binary operation
      (including those performed in parallel) is
      `2 * (N / 2 + N / 4 + ... + 1) = 2N - 2`
      --- i.e., approximately twice as many as a naive approach.
  
      [1] Blelloch, Guy E.
          [Prefix sums and their applications](
          https://www.cs.cmu.edu/~guyb/papers/Ble93.pdf)
          Technical Report CMU-CS-90-190,
          School of Computer Science,
          Carnegie Mellon University, 1990.
  
      Args:
        fn: Python callable implementing an associative binary operation with
          signature `r = fn(a, b)`. This must satisfy associativity:
          `fn(a, fn(b, c)) == fn(fn(a, b), c)`. The inputs and result are
          (possibly nested structures of) `Tensor`(s), matching `elems`. Each
          `Tensor` has a batch dimension in place of `elem_length`; the `fn`
          is expected to map over this dimension. The result `r` has the same shape
          (and structure) as the two inputs `a` and `b`.
        elems: A (possibly nested structure of) `Tensor`(s), each with dimension
          `elem_length` along `axis`. Note that `elem_length` determines the number
          of recursive steps required to perform the scan: if, in graph mode,
          this is not statically available, then ops will be created to
          handle any `elem_length` up to the maximum dimension of a `Tensor`.
        max_num_levels: Python `int`. The `axis` of the tensors in `elems` must have
          size less than `2**(max_num_levels + 1)`. The default value is
          sufficiently large for most needs. Lowering this value can reduce
          graph-building time when `scan_associative` is used with inputs of unknown
          shape.
          Default value: `48`.
        axis: Tensor `int` axis along which to perform the scan.
        validate_args: Python `bool`. When `True`, runtime checks
          for invalid inputs are performed. This may carry a performance cost.
          Default value: `False`.
        name: Python `str` name prefixed to ops created by this function.
      Returns:
        result: A (possibly nested structure of) `Tensor`(s) of the same shape
          and structure as `elems`, in which the `k`th element is the result of
          recursively applying `fn` to combine the first `k` elements of
          `elems`. For example, given `elems = [a, b, c, ...]`, the result
          would be `[a, fn(a, b), fn(fn(a, b), c), ...]`.
  
      #### Examples
  
      ```python
      import tensorflow as tf
      import tensorflow_probability as tfp
      import operator
  
      # Example 1: Partials sums of numbers.
  
      tfp.math.scan_associative(operator.add, tf.range(0, 4))
      # ==> [ 0, 1, 3, 6]
  
      # Example 2: Partial products of random matrices.
  
      dist = tfp.distributions.Normal(loc=0., scale=1.)
      matrices = dist.sample(sample_shape=[100, 2, 2])
      tfp.math.scan_associative(tf.matmul, matrices)
      ```
      """
      slice_elem = functools.partial(_slice_along_axis, axis=axis)
  
      def lowered_fn(a, b):
          # Lower `fn` to operate on flattened sequences of elems.
          with tf.name_scope("fn"):
              return tf.nest.flatten(
                  fn(
                      tf.nest.pack_sequence_as(elems, a),
                      tf.nest.pack_sequence_as(elems, b),
                  )
              )
  
      elems_flat = [
          tf.convert_to_tensor(elem) for elem in tf.nest.flatten(elems)
      ]
  
      # Summary of algorithm:
      #
      # Consider elements of `_scan(elems)` at odd indices. That's the same as first
      # summing successive pairs of elements of `elems` and performing a scan on
      # that half sized tensor. We perform the latter scan by recursion.
      #
      # Now consider the even elements of `_scan(elems)`. These can be computed
      # from the odd elements of `_scan(elems)` by adding each odd element of
      # `_scan(elems)` to the matching even element in the original `elems`.
      #
      # We return the odd and even elements interleaved.
      #
      # For the base case of the recursion we return the first element
      # of `elems` followed by the sum of the first two elements computed as
      # a (small two-down-to-one) reduction step.
  
      # The following is a pictorial representation of the algorithm using the
      # variables in the code below. The operator '+' is used to represent
      # the binary operation.
      # Note how the recursive call to `_scan` operates on a reduced form of the
      # input array in which successive pairs have already been summed.
  
      # elems     x0         x1   x2         x3    x4         x5    ...
      #           |\         /    | \        /     | \        /
      #           | \       /     |  \      /      |  \      /
      #           |  \     /      |   \    /       |   \    /
      #           |   \   /       |    \  /        |    \  /
      # reduced   |  x0+x1        |   x2+x3        |    x4+x5       ...
      # _elems    |    |          |     |          |       |
      #           |    |          |     |          |       |
      #           |    |          |     |          |       |
      # _scan(..) |    |          |     |          |       |
      #        +--|----+----------|-----+----------|-------+----    ...
      #        |  |               |                |
      #        |  |               |                |
      #        +--|----+----------|-----+----------|-------+----    ...
      #           |    |          |     |          |       |
      # odd       |  x0+x1        |   x0+...+x3    |     x0+..+x5   ...
      # _elems    |    | \        |     |      \   |       |
      #           |    |  \       |     |       \  |       |
      # even      |    |   \      |     |        \ |       |
      # _elems    x0   |   x0+...+x2    |       x0+...+x4  |        ...
      #           |    |          |     |          |       |
      # inter     |    |          |     |          |       |
      # leave(..) |    |          |     |          |       |
      #           x0 x0+x1 x0+...+x2  x0+...+x3 x0+...+x4 x0+...+x5 ...
  
      # TODO(b/150374456): if the sizes of all of the tensors can be determined
      # statically then we don't need a `level` parameter.
      def _scan(level, elems):
          """Perform scan on `elems`."""
          elem_length = ps.shape(elems[0])[axis]
  
          # Apply `fn` to reduce adjacent pairs to a single entry.
          a = [slice_elem(elem, 0, -1, step=2) for elem in elems]
          b = [slice_elem(elem, 1, None, step=2) for elem in elems]
          reduced_elems = lowered_fn(a, b)
  
          def handle_base_case_elem_length_two():
              return [
                  tf.concat(
                      [slice_elem(elem, 0, 1), reduced_elem], axis=axis
                  )
                  for (reduced_elem, elem) in zip(reduced_elems, elems)
              ]
  
          def handle_base_case_elem_length_three():
              reduced_reduced_elems = lowered_fn(
                  reduced_elems,
                  [slice_elem(elem, 2, 3) for elem in elems],
              )
              return [
                  tf.concat(
                      [
                          slice_elem(
                              elem, 0, 1
                          ),  # pylint: disable=g-complex-comprehension
                          reduced_elem,
                          reduced_reduced_elem,
                      ],
                      axis=axis,
                  )
                  for (reduced_reduced_elem, reduced_elem, elem) in zip(
                      reduced_reduced_elems, reduced_elems, elems
                  )
              ]
  
          # Base case of recursion: assumes `elem_length` is 2 or 3.
          at_base_case = ps.logical_or(
              ps.equal(elem_length, 2), ps.equal(elem_length, 3)
          )
          base_value = lambda: ps.cond(  # pylint: disable=g-long-lambda
              ps.equal(elem_length, 2),
              handle_base_case_elem_length_two,
              handle_base_case_elem_length_three,
          )
  
          if level <= 0:
              return base_value()
  
          def recursive_case():
              """Evaluate the next step of the recursion."""
              odd_elems = _scan(level - 1, reduced_elems)
  
              def even_length_case():
                  return lowered_fn(
                      [
                          slice_elem(odd_elem, 0, -1)
                          for odd_elem in odd_elems
                      ],
                      [slice_elem(elem, 2, None, 2) for elem in elems],
                  )
  
              def odd_length_case():
                  return lowered_fn(
                      [odd_elem for odd_elem in odd_elems],
                      [slice_elem(elem, 2, None, 2) for elem in elems],
                  )
  
              results = ps.cond(
                  ps.equal(elem_length % 2, 0),
                  even_length_case,
                  odd_length_case,
              )
  
              # The first element of a scan is the same as the first element
              # of the original `elems`.
              even_elems = [
                  tf.concat([slice_elem(elem, 0, 1), result], axis=axis)
                  for (elem, result) in zip(elems, results)
              ]
              return list(
                  map(
                      lambda a, b: _interleave(a, b, axis=axis),
                      even_elems,
                      odd_elems,
                  )
              )
  
          return ps.cond(at_base_case, base_value, recursive_case)
  
      with tf.name_scope(name if name else "scan_associative"):
          elem_length, assertions = _validate_elem_length(
              max_num_levels, elems_flat, axis=axis
          )
  
      with tf.control_dependencies(assertions if validate_args else []):
          return ps.cond(
              elem_length < 2,
              lambda: elems,
              lambda: (
                  tf.nest.pack_sequence_as(  # pylint: disable=g-long-lambda
                      elems, _scan(max_num_levels - 1, elems_flat)
                  )
              ),
          )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 4 for calling a global function with example usage: scan_associative(fn, elems, max_num_levels, axis, validate_args, name) and returns: ps, ps, tf, unknown, unknown, base_value, list, lowered_fn, lowered_fn