code: def test_cumsum_vector(self, total_steps=15):
      tree_aggregator = tree_aggregation.TreeAggregator(
          value_generator=ConstantValueGenerator(
              [
                  tf.ones([2, 2], dtype=tf.float32),
                  tf.constant([2], dtype=tf.float32),
              ]
          )
      )
      tree_aggregator_truth = tree_aggregation.TreeAggregator(
          value_generator=ConstantValueGenerator(1.0)
      )
      state = tree_aggregator.init_state()
      truth_state = tree_aggregator_truth.init_state()
      for leaf_node_idx in range(total_steps):
          self.assertEqual(
              leaf_node_idx, tree_aggregation.get_step_idx(state)
          )
          val, state = tree_aggregator.get_cumsum_and_update(state)
          (
              expected_val,
              truth_state,
          ) = tree_aggregator_truth.get_cumsum_and_update(truth_state)
          self.assertEqual(
              tree_aggregation.get_step_idx(state),
              tree_aggregation.get_step_idx(truth_state),
          )
          expected_result = [
              expected_val * tf.ones([2, 2], dtype=tf.float32),
              expected_val * tf.constant([2], dtype=tf.float32),
          ]
          tf.nest.map_structure(
              self.assertAllEqual, val, expected_result
          )
  
imports: import math
  import random
  from absl.testing import parameterized
  import tensorflow
  from tensorflow_privacy.privacy.dp_query import tree_aggregation
label_desc: Use image node_id 6 for calling the TreeAggregatorTest obj's underlying member method code with example usage: obj.test_cumsum_vector(total_steps) without return types