code: def _queue_push(queue, should_update, new_vecs):
      """Conditionally push new vectors into a batch of first-in-first-out queues.
  
      The `queue` of shape `[k, ..., n]` can be thought of as a batch of queues,
      each holding `k` n-D vectors; while `new_vecs` of shape `[..., n]` is a
      fresh new batch of n-D vectors. The `should_update` batch of Boolean scalars,
      i.e. shape `[...]`, indicates batch members whose corresponding n-D vector in
      `new_vecs` should be added at the back of its queue, pushing out the
      corresponding n-D vector from the front. Batch members in `new_vecs` for
      which `should_update` is False are ignored.
  
      Note: the choice of placing `k` at the dimension 0 of the queue is
      constrained by the L-BFGS two-loop algorithm above. The algorithm uses
      tf.scan to iterate over the `k` correction pairs simulatneously across all
      batches, and tf.scan itself can only iterate over dimension 0.
  
      For example:
  
      ```python
        k, b, n = (3, 2, 5)
        queue = tf.reshape(tf.range(30), (k, b, n))
        # => [[[ 0,  1,  2,  3,  4],
        #      [ 5,  6,  7,  8,  9]],
        #
        #     [[10, 11, 12, 13, 14],
        #      [15, 16, 17, 18, 19]],
        #
        #     [[20, 21, 22, 23, 24],
        #      [25, 26, 27, 28, 29]]]
  
        element = tf.reshape(tf.range(30, 40), (b, n))
        # => [[30, 31, 32, 33, 34],
              [35, 36, 37, 38, 39]]
  
        should_update = tf.constant([True, False])  # Shape: (b,)
  
        _queue_add(should_update, queue, element)
        # => [[[10, 11, 12, 13, 14],
        #      [ 5,  6,  7,  8,  9]],
        #
        #     [[20, 21, 22, 23, 24],
        #      [15, 16, 17, 18, 19]],
        #
        #     [[30, 31, 32, 33, 34],
        #      [25, 26, 27, 28, 29]]]
      ```
  
      Args:
        queue: A `tf.Tensor` of shape `[k, ..., n]`; a batch of queues each with
          `k` n-D vectors.
        should_update: A Boolean `tf.Tensor` of shape `[...]` indicating batch
          members where new vectors should be added to their queues.
        new_vecs: A `tf.Tensor` of shape `[..., n]`; a batch of n-D vectors to add
          at the end of their respective queues, pushing out the first element from
          each.
  
      Returns:
        A new `tf.Tensor` of shape `[k, ..., n]`.
      """
      new_queue = tf.concat([queue[1:], [new_vecs]], axis=0)
      return tf.where(
          should_update[tf.newaxis, ..., tf.newaxis], new_queue, queue
      )
  
file_path: probability/tensorflow_probability/python/optimizer/lbfgs.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.optimizer import bfgs_utils
label_desc: Use image node_id 5 for calling a global function with example usage: _queue_push(queue, should_update, new_vecs) and returns: tf