code: def _build_tree_from_leaf(
      leaf_nodes: tf.Tensor, arity: int
  ) -> tf.RaggedTensor:
      """A function constructs a complete tree given all the leaf nodes.
  
      The function takes a 1-D array representing the leaf nodes of a tree and the
      tree's arity, and constructs a complete tree by recursively summing the
      adjacent children to get the parent until reaching the root node. Because we
      assume a complete tree, if the number of leaf nodes does not divide arity, the
      leaf nodes will be padded with zeros.
  
      Args:
        leaf_nodes: A 1-D array storing the leaf nodes of the tree.
        arity: A `int` for the branching factor of the tree, i.e. the number of
          children for each internal node.
  
      Returns:
        `tf.RaggedTensor` representing the tree. For example, if
        `leaf_nodes=tf.Tensor([1, 2, 3, 4])` and `arity=2`, then the returned value
        should be `tree=tf.RaggedTensor([[10],[3,7],[1,2,3,4]])`. In this way,
        `tree[layer][index]` can be used to access the node indexed by (layer,
        index) in the tree,
      """
  
      def pad_zero(leaf_nodes, size):
          paddings = tf.zeros(
              shape=(size - leaf_nodes.shape[0],),
              dtype=leaf_nodes.dtype,
          )
          return tf.concat((leaf_nodes, paddings), axis=0)
  
      leaf_nodes_size = tf.constant(
          leaf_nodes.shape[0], dtype=tf.float32
      )
      num_layers = (
          tf.math.ceil(
              tf.math.log(leaf_nodes_size)
              / tf.math.log(tf.cast(arity, dtype=tf.float32))
          )
          + 1
      )
      leaf_nodes = pad_zero(
          leaf_nodes,
          tf.math.pow(tf.cast(arity, dtype=tf.float32), num_layers - 1),
      )
  
      def _shrink_layer(layer: tf.Tensor, arity: int) -> tf.Tensor:
          return tf.reduce_sum((tf.reshape(layer, (-1, arity))), 1)
  
      # The following `tf.while_loop` constructs the tree from bottom up by
      # iteratively applying `_shrink_layer` to each layer of the tree. The reason
      # for the choice of TF1.0-style `tf.while_loop` is that @tf.function does not
      # support auto-translation from python loop to tf loop when loop variables
      # contain a `RaggedTensor` whose shape changes across iterations.
  
      idx = tf.identity(num_layers)
      loop_cond = lambda i, h: tf.less_equal(2.0, i)
  
      def _loop_body(i, h):
          return [
              tf.add(i, -1.0),
              tf.concat(([_shrink_layer(h[0], arity)], h), axis=0),
          ]
  
      _, tree = tf.while_loop(
          loop_cond,
          _loop_body,
          [idx, tf.RaggedTensor.from_tensor([leaf_nodes])],
          shape_invariants=[
              idx.get_shape(),
              tf.RaggedTensorSpec(
                  dtype=leaf_nodes.dtype, ragged_rank=1
              ),
          ],
      )
  
      return tree
  
file_path: privacy/tensorflow_privacy/privacy/dp_query/tree_range_query.py
imports: import distutils
  import math
  from typing import Any, NamedTuple, Optional
  import dp_accounting
  import tensorflow
  from tensorflow_privacy.privacy.dp_query import distributed_discrete_gaussian_query
  from tensorflow_privacy.privacy.dp_query import dp_query
  from tensorflow_privacy.privacy.dp_query import gaussian_query
label_desc: Use image node_id 1 for calling a global function with example usage: _build_tree_from_leaf(leaf_nodes, arity) and returns: tree, tf, tf, list