class_docstr: Immutable, internal container for a fixed size stack.
  
  The implementation is backed by a `Tensor` each for the stack and the
  (batched) stack pointer.
  
  As a namedtuple, it can be directly passed through TF's nest library for
  flattening and restructuring as an element passed to e.g. a TF while loop.
docstr: Pushes `value` onto the stack, advances frame of batch members in `mask`.
  
  In this impl, we update each thread's top-of-stack (regardless of `mask`) to
  the corresponding `value`, then advance the stack pointers of only those
  threads indicated by `mask`.
  
  Args:
    value: `Tensor` having the shape of a single batch of the variable.
    mask: Boolean `Tensor` of shape `[batch_size]`. Threads at `True` indices
        of `mask` have their stack frames advanced; the others remain.
    name: Optional name for this op.
  
  Returns:
    stack: Updated stack. Does not mutate `self`.
    asserted_value: A assertion-bound snapshot of the input `value`,
        assertions used to catch stack overflows.
docstr_tok: ["Pushes","`","value","`","onto","the","stack",",","advances","frame","of","batch","members","in","`","mask","`",".","In","this","impl",",","we","update","each","thread","'s","top-of-stack","(","regardless","of","`","mask","`",")","to","the","corresponding","`","value","`",",","then","advance","the","stack","pointers","of","only","those","threads","indicated","by","`","mask","`",".","Args",":","value",":","`","Tensor","`","having","the","shape","of","a","single","batch","of","the","variable",".","mask",":","Boolean","`","Tensor","`","of","shape","`","[","batch_size","]","`",".","Threads","at","`","True","`","indices","of","`","mask","`","have","their","stack","frames","advanced",";","the","others","remain",".","name",":","Optional","name","for","this","op",".","Returns",":","stack",":","Updated","stack",".","Does","not","mutate","`","self","`",".","asserted_value",":","A","assertion-bound","snapshot","of","the","input","`","value","`",",","assertions","used","to","catch","stack","overflows","."]
code: def push(self, value, mask, name=None):
      """Pushes `value` onto the stack, advances frame of batch members in `mask`.
  
      In this impl, we update each thread's top-of-stack (regardless of `mask`) to
      the corresponding `value`, then advance the stack pointers of only those
      threads indicated by `mask`.
  
      Args:
        value: `Tensor` having the shape of a single batch of the variable.
        mask: Boolean `Tensor` of shape `[batch_size]`. Threads at `True` indices
            of `mask` have their stack frames advanced; the others remain.
        name: Optional name for this op.
  
      Returns:
        stack: Updated stack. Does not mutate `self`.
        asserted_value: A assertion-bound snapshot of the input `value`,
            assertions used to catch stack overflows.
      """
      with tf.name_scope(name or "Stack.push"):
          value = tf.convert_to_tensor(value=value, name="value")
          mask = tf.convert_to_tensor(value=mask, name="mask")
          # self.stack:       [max_stack_depth * batch_size, ...]
          # self.stack_index:                   [batch_size]
          # value:                              [batch_size, ...]
          batch_size = (
              tf.compat.dimension_value(self.stack_index.shape[0])
              or tf.shape(input=self.stack_index)[0]
          )
          max_stack_depth = (
              tf.compat.dimension_value(self.stack.shape[0])
              or tf.shape(input=self.stack)[0]
          ) // batch_size
          max_stack_depth_tensor = tf.convert_to_tensor(
              value=max_stack_depth
          )
          tiled_value = tf.tile(
              input=value[tf.newaxis, ...],
              multiples=tf.concat(
                  [
                      [max_stack_depth_tensor],
                      tf.ones(
                          tf.rank(value),
                          dtype=max_stack_depth_tensor.dtype,
                      ),
                  ],
                  axis=0,
              ),
          )
          update_stack_mask = tf.one_hot(
              self.stack_index,
              depth=max_stack_depth,
              axis=0,  # Stack depth x batch are both in outermost dim, stack major.
              on_value=True,
              off_value=False,
              dtype=tf.bool,
          )
          new_stack = tf1.where(
              tf.reshape(update_stack_mask, [-1]),
              tf.reshape(tiled_value, tf.shape(input=self.stack)),
              self.stack,
          )
          new_stack.set_shape(self.stack.shape)
          new_stack_index = self.stack_index + tf.cast(
              mask, self.stack_index.dtype
          )
          new_stack_index.set_shape(self.stack_index.shape)
          if self._safety_checks():
              with tf.control_dependencies(
                  [
                      tf1.assert_less(
                          new_stack_index,
                          tf.cast(
                              max_stack_depth_tensor,
                              new_stack_index.dtype,
                          ),
                      )
                  ]
              ):
                  value = tf.identity(value)
                  new_stack_index = tf.identity(new_stack_index)
          return type(self)(new_stack, new_stack_index), value
  
imports: import collections
  import contextlib
  import six
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.auto_batching import instructions
  from tensorflow_probability.python.experimental.auto_batching import xla
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 3 for calling the Stack obj's underlying member method code with example usage: obj.push(value, mask, name) and returns: type,  value