class_docstr: A bounded functional stack implementation. Elements may be pytrees.
docstr: Creates a stack with size `capacity` with elements like `prototype`.
  
  `prototype` can be any JAX pytree. This function looks only at its
  structure; the specific values are ignored.
docstr_tok: ["Creates","a","stack","with","size","`","capacity","`","with","elements","like","`","prototype","`",".","`","prototype","`","can","be","any","JAX","pytree",".","This","function","looks","only","at","its","structure",";","the","specific","values","are","ignored","."]
code: def create(capacity: int, prototype: Any) -> Stack:
      """Creates a stack with size `capacity` with elements like `prototype`.
  
      `prototype` can be any JAX pytree. This function looks only at its
      structure; the specific values are ignored.
      """
      return Stack(
          jnp.array(0, jnp.int32),
          jax.tree_util.tree_map(
              lambda x: jnp.zeros(
                  (capacity,) + tuple(x.shape), x.dtype
              ),
              prototype,
          ),
      )
  
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 3 for calling the Stack obj's underlying member method code with example usage: obj.create(capacity, prototype) and returns: Stack