code: def test_tree_sum_noise_expected(
      self,
      total_steps,
      expected_variance,
      noise_std,
      variable_shape,
      tolerance,
  ):
      # Test whether `tree_aggregator` will output `expected_variance` (within a
      # relative `tolerance`) after  `total_steps` of leaf nodes are traversed.
      # Each tree node is a `variable_shape` tensor of Gaussian noise with
      # `noise_std`. Note that the variance of a tree node is smaller than
      # the given vanilla node `noise_std` because of the update rule of
      # `EfficientTreeAggregator`.
      random_generator = tree_aggregation.GaussianNoiseGenerator(
          noise_std, tf.TensorSpec(variable_shape), seed=2020
      )
      tree_aggregator = tree_aggregation.EfficientTreeAggregator(
          value_generator=random_generator
      )
      state = tree_aggregator.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)
      self.assertAllClose(
          math.sqrt(expected_variance),
          tf.math.reduce_std(val),
          rtol=tolerance,
      )
  
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 2 for calling the EfficientTreeAggregatorTest obj's underlying member method code with example usage: obj.test_tree_sum_noise_expected(total_steps, expected_variance, noise_std, variable_shape, tolerance) without return types