class_docstr: A bounded functional stack implementation. Elements may be pytrees.
docstr: Pushes `elem` onto the stack, returning the updated stack.
docstr_tok: ["Pushes","`","elem","`","onto","the","stack",",","returning","the","updated","stack","."]
code: def push(self, elem: Any) -> Stack:
      """Pushes `elem` onto the stack, returning the updated stack."""
      return Stack(
          self._size + 1,
          jax.tree_util.tree_map(
              lambda x, y: lax.dynamic_update_index_in_dim(
                  x, y, self._size, 0
              ),
              self._data,
              elem,
          ),
      )
  
file_path: jax/jax/_src/lax/stack.py
imports: from __future__ import annotations
  from typing import Any
  import jax
  from jax import lax
  import jax.numpy
label_desc: Use image node_id 5 for calling the Stack obj's underlying member method code with example usage: obj.push(elem) and returns: Stack