code: def test_counting_by_matmul_example(self):
      num_elems = 2**4 + 1
      upper_row = tf.stack(
          [
              tf.ones(num_elems, dtype=tf.int64),
              tf.range(0, num_elems, dtype=tf.int64),
          ],
          axis=1,
      )
      lower_row = tf.stack(
          [
              tf.zeros(num_elems, dtype=tf.int64),
              tf.ones(num_elems, dtype=tf.int64),
          ],
          axis=1,
      )
      m = self._maybe_static(tf.stack([upper_row, lower_row], axis=1))
      result = self.evaluate(
          scan_associative(tf.matmul, m, max_num_levels=8)
      )
      self.assertAllEqual(
          result[:, 0, 1], np.cumsum(np.arange(num_elems))
      )
  
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 6 for calling the _ScanAssociativeTest obj's underlying member method code with example usage: obj.test_counting_by_matmul_example() without return types