xtructure.stack package
Submodules
xtructure.stack.stack module
- class xtructure.stack.stack.Stack(max_size: int, size: uint32, val_store: Xtructurable)[source]
Bases:
objectA JAX-compatible batched Stack data structure. Optimized for parallel operations on GPU using JAX.
- max_size
Maximum number of elements the stack can hold.
- Type:
int
- size
Current number of elements in the stack.
- Type:
jax.numpy.uint32
- val_store
Array storing the values in the stack.
- static build(max_size: int, value_class: Xtructurable)[source]
Creates a new Stack instance.
- Parameters:
max_size – The maximum number of elements the stack can hold.
value_class – The class of values to be stored in the stack. It must be a subclass of Xtructurable.
- Returns:
A new, empty Stack instance.
- from_tuple()
- max_size: int
- peek(num_items: int = 1)[source]
Peeks at the top items of the stack without removing them.
- Parameters:
num_items – The number of items to peek at. Defaults to 1.
- Returns:
The top num_items from the stack.
- pop(num_items: int = 1)[source]
Pops a number of items from the stack.
- Parameters:
num_items – The number of items to pop.
- Returns:
A new Stack instance with items removed.
The popped items.
- Return type:
A tuple containing
- push(items: Xtructurable)[source]
Pushes a batch of items onto the stack.
- Parameters:
items – An Xtructurable containing the items to push. The first dimension is the batch dimension.
- Returns:
A new Stack instance with the items pushed onto it.
- replace(**kwargs)
- size: uint32
- to_tuple()
- val_store: Xtructurable
Module contents
- class xtructure.stack.Stack(max_size: int, size: uint32, val_store: Xtructurable)[source]
Bases:
objectA JAX-compatible batched Stack data structure. Optimized for parallel operations on GPU using JAX.
- max_size
Maximum number of elements the stack can hold.
- Type:
int
- size
Current number of elements in the stack.
- Type:
jax.numpy.uint32
- val_store
Array storing the values in the stack.
- static build(max_size: int, value_class: Xtructurable)[source]
Creates a new Stack instance.
- Parameters:
max_size – The maximum number of elements the stack can hold.
value_class – The class of values to be stored in the stack. It must be a subclass of Xtructurable.
- Returns:
A new, empty Stack instance.
- from_tuple()
- max_size: int
- peek(num_items: int = 1)[source]
Peeks at the top items of the stack without removing them.
- Parameters:
num_items – The number of items to peek at. Defaults to 1.
- Returns:
The top num_items from the stack.
- pop(num_items: int = 1)[source]
Pops a number of items from the stack.
- Parameters:
num_items – The number of items to pop.
- Returns:
A new Stack instance with items removed.
The popped items.
- Return type:
A tuple containing
- push(items: Xtructurable)[source]
Pushes a batch of items onto the stack.
- Parameters:
items – An Xtructurable containing the items to push. The first dimension is the batch dimension.
- Returns:
A new Stack instance with the items pushed onto it.
- replace(**kwargs)
- size: uint32
- to_tuple()
- val_store: Xtructurable