code: def test_cumulative_sum_custom_axis(self, axis):
      elems = self._maybe_static(
          tf.random.stateless_normal(
              [4, 32, 31, 1],
              seed=test_util.test_seed(sampler_type="stateless"),
          )
      )
  
      axis = self._maybe_static(axis)
      expected_result = self.evaluate(tf.cumsum(elems, axis=axis))
      result = scan_associative(
          operator.add, elems, axis=axis, max_num_levels=8
      )
      self.assertAllClose(
          self.evaluate(result), expected_result, rtol=1e-5
      )
  
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 5 for calling the _ScanAssociativeTest obj's underlying member method code with example usage: obj.test_cumulative_sum_custom_axis(axis) without return types