code: def test_cumulative_sum_with_xla(self):
      elems = self._maybe_static(
          tf.range(0, 2**4 - 1, dtype=tf.int64)
      )
  
      # JAX jit expects arguments to functions to be DeviceArrays. Thus we
      # curry `scan_associative` so that the resulting function takes in `Tensors`
      # or `DeviceArrays`.
  
      xla_scan = tf.function(jit_compile=True)(
          functools.partial(scan_associative, operator.add)
      )
      result = xla_scan(elems)
  
      self.assertAllEqual(
          self.evaluate(result), self.evaluate(tf.cumsum(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 1 for calling the ScanAssociativeTestStatic obj's underlying member method code with example usage: obj.test_cumulative_sum_with_xla() without return types