code: def get_candidate_values(
      x: tf.Tensor,  # (batch_size, ...)
      candidate_ids: tf.Tensor,  # (batch_size, num_candidates)
  ) -> tf.Tensor:
      """Gathers candidate values according to IDs.
  
      Args:
          x: Any tensor with at least one dimension
          candidate_ids: Indicator for which candidates to gather
  
      Returns:
          A tensor of shape `(batch_size, 1, num_candidates, tf.shape(x)[-1])`, where
          for each batch example, we generate a list of `num_candidates` vectors, and
          each candidate is chosen from `x` according to the candidate id. For example:
  
          ```
          x = [[0 1 2],
                  [3 4 5],
                  [6 7 8]]
          candidate_ids = [[0, 1], [0, 0], [2, 0]]
          gives
          [
              [[0 1 2],
               [3 4 5]],
              [[0 1 2],
               [0 1 2]],
              [[6 7 8],
               [0 1 2]]
          ]
          ```
      """
      tiled_x = tf.tile(
          tf.expand_dims(batch_flatten(x), 0),
          (tf.shape(candidate_ids)[0], 1, 1),
      )
      candidate_values = tf.gather(tiled_x, candidate_ids, batch_dims=1)
  
      return candidate_values
  
imports: import tensorflow
  from tensorflow import Tensor
  from typing import Union, Optional
label_desc: Use image node_id 3 for calling a global function with example usage: get_candidate_values(x, candidate_ids) and returns: candidate_values