class_docstr: Internal container for a batched stack.
  
  The implementation is a preallocated array and a (batched) stack
  pointer.
  
  The namedtuple structure exposes the full state of the stack, and is useful
  for testing, passing through flatten/unflatten operations, and general
  symmetry with the TensorFlow backend.
class_docstr_tok: ["Internal","container","for","a","batched","stack",".","The","implementation","is","a","preallocated","array","and","a","(","batched",")","stack","pointer",".","The","namedtuple","structure","exposes","the","full","state","of","the","stack",",","and","is","useful","for","testing",",","passing","through","flatten\/unflatten","operations",",","and","general","symmetry","with","the","TensorFlow","backend","."]
docstr: Writes new value to all threads, updates frame of those in `mask`.
  
  Args:
    value: Value to write into all threads top frame before updating `mask`
      frame pointers.
    mask: Boolean array of shape `[batch_size]`. The values at `True`
      indices of `mask` are updated; the others remain.
  
  Returns:
    stack: Updated stack. Does not mutate `self`.
  
  Raises:
    ValueError: If a push exceeds the maximum stack depth.
docstr_tok: ["Writes","new","value","to","all","threads",",","updates","frame","of","those","in","`","mask","`",".","Args",":","value",":","Value","to","write","into","all","threads","top","frame","before","updating","`","mask","`","frame","pointers",".","mask",":","Boolean","array","of","shape","`","[","batch_size","]","`",".","The","values","at","`","True","`","indices","of","`","mask","`","are","updated",";","the","others","remain",".","Returns",":","stack",":","Updated","stack",".","Does","not","mutate","`","self","`",".","Raises",":","ValueError",":","If","a","push","exceeds","the","maximum","stack","depth","."]
code: def push(self, value, mask):
      """Writes new value to all threads, updates frame of those in `mask`.
  
      Args:
        value: Value to write into all threads top frame before updating `mask`
          frame pointers.
        mask: Boolean array of shape `[batch_size]`. The values at `True`
          indices of `mask` are updated; the others remain.
  
      Returns:
        stack: Updated stack. Does not mutate `self`.
  
      Raises:
        ValueError: If a push exceeds the maximum stack depth.
      """
      # stack:       [max_stack_depth * batch_size, ...]
      # stack_index:                   [batch_size]
      # value:                         [batch_size, ...]
      batch_size = self.stack_index.shape[0]
      max_stack_depth = self.stack.shape[0] // batch_size
      tiled_value = np.reshape(
          np.repeat(value[np.newaxis], max_stack_depth, axis=0),
          self.stack.shape,
      )
      update_indices = self.stack_index * batch_size + np.arange(
          batch_size
      )
      new_stack = _where_leading_dims(
          np.any(
              np.equal(
                  update_indices,
                  np.arange(self.stack.shape[0])[:, np.newaxis],
              ),
              axis=1,
          ),
          tiled_value,
          self.stack,
      )
      new_stack_index = self.stack_index + mask
      if np.any(new_stack_index >= max_stack_depth):
          raise ValueError("Maximum stack depth exceeded.")
      return Stack(new_stack, new_stack_index)
  
imports: import collections
  import numpy
  from tensorflow_probability.python.experimental.auto_batching import instructions
label_desc: Use image node_id 2 for calling the Stack obj's underlying member method code with example usage: obj.push(value, mask) and returns: Stack