code: def test_can_scan_tensors_of_different_rank(self):
      num_elems = 2**4
      elems0 = self.evaluate(
          uniform.Uniform(-1.0, 1.0).sample(
              sample_shape=[num_elems], seed=test_util.test_seed()
          )
      )
      elems1 = self.evaluate(
          uniform.Uniform(-1.0, 1.0).sample(
              sample_shape=[num_elems, 1], seed=test_util.test_seed()
          )
      )
  
      def extended_add(a, b):
          return (a[0] + b[0], a[1] + b[1])
  
      result = self.evaluate(
          scan_associative(
              extended_add,
              (self._maybe_static(elems0), self._maybe_static(elems1)),
              max_num_levels=8,
          )
      )
  
      self.assertAllClose(result[0], self.evaluate(tf.cumsum(elems0)))
      self.assertAllClose(
          result[1], self.evaluate(tf.cumsum(elems1, axis=0))
      )
  
imports: import collections
  import functools
  import operator
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.math.scan_associative import scan_associative
label_desc: Use image node_id 9 for calling the _ScanAssociativeTest obj's underlying member method code with example usage: obj.test_can_scan_tensors_of_different_rank() and returns: unknown,  unknown