tensorflow_privacy/privacy/dp_query/tree_aggregation.py [241:274]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    level_buffer_idx = tf.TensorArray(dtype=tf.int32, size=1, dynamic_size=True)
    level_buffer_idx = level_buffer_idx.write(0, tf.constant(
        0, dtype=tf.int32)).stack()

    new_val, value_generator_state = self.value_generator.next(
        value_generator_state)
    level_buffer_structure = tf.nest.map_structure(
        lambda x: tf.TensorArray(dtype=tf.float32, size=1, dynamic_size=True),
        new_val)
    level_buffer = tf.nest.map_structure(lambda x, y: x.write(0, y).stack(),
                                         level_buffer_structure, new_val)
    return TreeState(
        level_buffer=level_buffer,
        level_buffer_idx=level_buffer_idx,
        value_generator_state=value_generator_state)

  def init_state(self) -> TreeState:
    """Returns initial `TreeState`.

    Initializes `TreeState` for a tree of a single leaf node: the respective
    initial node value in `TreeState.level_buffer` is generated by the value
    generator function, and the node index is 0.

    Returns:
      An initialized `TreeState`.
    """
    value_generator_state = self.value_generator.initialize()
    return self._get_init_state(value_generator_state)

  def reset_state(self, state: TreeState) -> TreeState:
    """Returns reset `TreeState` after restarting a new tree."""
    return self._get_init_state(state.value_generator_state)

  @tf.function
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



tensorflow_privacy/privacy/dp_query/tree_aggregation.py [392:425]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    level_buffer_idx = tf.TensorArray(dtype=tf.int32, size=1, dynamic_size=True)
    level_buffer_idx = level_buffer_idx.write(0, tf.constant(
        0, dtype=tf.int32)).stack()

    new_val, value_generator_state = self.value_generator.next(
        value_generator_state)
    level_buffer_structure = tf.nest.map_structure(
        lambda x: tf.TensorArray(dtype=tf.float32, size=1, dynamic_size=True),
        new_val)
    level_buffer = tf.nest.map_structure(lambda x, y: x.write(0, y).stack(),
                                         level_buffer_structure, new_val)
    return TreeState(
        level_buffer=level_buffer,
        level_buffer_idx=level_buffer_idx,
        value_generator_state=value_generator_state)

  def init_state(self) -> TreeState:
    """Returns initial `TreeState`.

    Initializes `TreeState` for a tree of a single leaf node: the respective
    initial node value in `TreeState.level_buffer` is generated by the value
    generator function, and the node index is 0.

    Returns:
      An initialized `TreeState`.
    """
    value_generator_state = self.value_generator.initialize()
    return self._get_init_state(value_generator_state)

  def reset_state(self, state: TreeState) -> TreeState:
    """Returns reset `TreeState` after restarting a new tree."""
    return self._get_init_state(state.value_generator_state)

  @tf.function
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -