xtructure.bgpq package

Subpackages

• xtructure.bgpq.merge_split package
  • Submodules
  • xtructure.bgpq.merge_split.common module
    • binary_search_partition()
  • xtructure.bgpq.merge_split.loop module
    • merge_arrays_indices_loop()
    • merge_indices_kernel_loop()
  • xtructure.bgpq.merge_split.parallel module
    • merge_arrays_parallel()
    • merge_parallel_kernel()
  • xtructure.bgpq.merge_split.split module
    • merge_sort_split_idx()
  • Module contents
    • merge_arrays_indices_loop()
    • merge_arrays_parallel()
    • merge_sort_split_idx()

Submodules

xtructure.bgpq.bgpq module

Batched GPU Priority Queue (BGPQ) Implementation This module provides a JAX-compatible priority queue optimized for GPU operations. Key features: - Fully batched operations for GPU efficiency - Supports custom value types through dataclass - Uses infinity padding for unused slots - Maintains sorted order for efficient min/max operations

class xtructure.bgpq.bgpq.BGPQ(max_size: int, heap_size: int, buffer_size: int, branch_size: int, batch_size: int, key_store: Array | ndarray | bool | number, val_store: Xtructurable, key_buffer: Array | ndarray | bool | number, val_buffer: Xtructurable)

Bases: object

Batched GPU Priority Queue implementation. Optimized for parallel operations on GPU using JAX.

max_size

Maximum number of elements the queue can hold

Type:

int

size

Current number of elements in the queue

branch_size

Number of branches in the heap tree

Type:

int

batch_size

Size of batched operations

Type:

int

key_store

Array storing keys in a binary heap structure

Type:

jax.Array | numpy.ndarray | numpy.bool | numpy.number

val_store

Array storing associated values

Type:

xtructure.core.protocol.Xtructurable

key_buffer

Buffer for keys waiting to be inserted

Type:

jax.Array | numpy.ndarray | numpy.bool | numpy.number

val_buffer

Buffer for values waiting to be inserted

Type:

xtructure.core.protocol.Xtructurable

batch_size: int

branch_size: int

buffer_size: int

static build(total_size, batch_size, value_class=<class 'xtructure.core.protocol.Xtructurable'>, key_dtype=<class 'jax.numpy.float16'>)

Create a new BGPQ instance with specified capacity.

Parameters:

• total_size – Total number of elements the queue can store

• batch_size – Size of batched operations

• value_class – Class to use for storing values (must implement default())

Returns:

A new priority queue instance initialized with empty storage

Return type:

BGPQ

static delete_heapify(heap: BGPQ)

Maintain heap property after deletion of minimum elements.

Parameters:

heap – The priority queue instance

Returns:

Updated heap instance

delete_mins()

Remove and return the minimum elements from the queue.

Parameters:

heap – The priority queue instance

Returns:

• Updated heap instance

• Array of minimum keys removed

• Xtructurable of corresponding values

Return type:

tuple containing

from_tuple()

heap_size: int

insert(block_key: Array | ndarray | bool | number, block_val: Xtructurable)

Insert new elements into the priority queue. Maintains heap property through merge operations and heapification.

Parameters:

• heap – The priority queue instance

• block_key – Keys to insert

• block_val – Values to insert

• added_size – Optional size of insertion (calculated if None)

Returns:

Updated heap instance

key_buffer: Array | ndarray | bool | number

key_store: Array | ndarray | bool | number

static make_batched(key: Array | ndarray | bool | number, val: Xtructurable, batch_size: int)

Convert unbatched arrays into batched format suitable for the queue.

Parameters:

• key – Array of keys to batch

• val – Xtructurable of values to batch

• batch_size – Desired batch size

Returns:

• Batched key array

• Batched value array

Return type:

tuple containing

max_size: int

merge_buffer(blockk: Array | ndarray | bool | number, blockv: Xtructurable)

Merge buffer contents with block contents, handling overflow conditions.

This method is crucial for maintaining the heap property when inserting new elements. It handles the case where the buffer might overflow into the main storage.

Parameters:

• blockk – Block keys array

• blockv – Block values

• bufferk – Buffer keys array

• bufferv – Buffer values

Returns:

• Updated block keys

• Updated block values

• Updated buffer keys

• Updated buffer values

• Boolean indicating if buffer overflow occurred

Return type:

tuple containing

replace(**kwargs)

property size

to_tuple()

val_buffer: Xtructurable

val_store: Xtructurable

xtructure.bgpq.bgpq.merge_sort_split(ak: Array | ndarray | bool | number, av: Xtructurable, bk: Array | ndarray | bool | number, bv: Xtructurable) → tuple[Array | ndarray | bool | number, Xtructurable, Array | ndarray | bool | number, Xtructurable]

Merge and split two sorted arrays while maintaining their relative order. This is a key operation for maintaining heap property in batched operations.

Parameters:

• ak – First array of keys

• av – First array of values

• bk – Second array of keys

• bv – Second array of values

Returns:

• First half of merged and sorted keys

• First half of corresponding values

• Second half of merged and sorted keys

• Second half of corresponding values

Return type:

tuple containing

xtructure.bgpq.bgpq.sort_arrays(k: Array | ndarray | bool | number, v: Xtructurable)

Module contents

class xtructure.bgpq.BGPQ(max_size: int, heap_size: int, buffer_size: int, branch_size: int, batch_size: int, key_store: Array | ndarray | bool | number, val_store: Xtructurable, key_buffer: Array | ndarray | bool | number, val_buffer: Xtructurable)

Bases: object

Batched GPU Priority Queue implementation. Optimized for parallel operations on GPU using JAX.

max_size

Maximum number of elements the queue can hold

Type:

int

size

Current number of elements in the queue

branch_size

Number of branches in the heap tree

Type:

int

batch_size

Size of batched operations

Type:

int

key_store

Array storing keys in a binary heap structure

Type:

jax.Array | numpy.ndarray | numpy.bool | numpy.number

val_store

Array storing associated values

Type:

xtructure.core.protocol.Xtructurable

key_buffer

Buffer for keys waiting to be inserted

Type:

jax.Array | numpy.ndarray | numpy.bool | numpy.number

val_buffer

Buffer for values waiting to be inserted

Type:

xtructure.core.protocol.Xtructurable

batch_size: int

branch_size: int

buffer_size: int

static build(total_size, batch_size, value_class=<class 'xtructure.core.protocol.Xtructurable'>, key_dtype=<class 'jax.numpy.float16'>)

Create a new BGPQ instance with specified capacity.

Parameters:

• total_size – Total number of elements the queue can store

• batch_size – Size of batched operations

• value_class – Class to use for storing values (must implement default())

Returns:

A new priority queue instance initialized with empty storage

Return type:

BGPQ

static delete_heapify(heap: BGPQ)

Maintain heap property after deletion of minimum elements.

Parameters:

heap – The priority queue instance

Returns:

Updated heap instance

delete_mins()

Remove and return the minimum elements from the queue.

Parameters:

heap – The priority queue instance

Returns:

• Updated heap instance

• Array of minimum keys removed

• Xtructurable of corresponding values

Return type:

tuple containing

from_tuple()

heap_size: int

insert(block_key: Array | ndarray | bool | number, block_val: Xtructurable)

Insert new elements into the priority queue. Maintains heap property through merge operations and heapification.

Parameters:

• heap – The priority queue instance

• block_key – Keys to insert

• block_val – Values to insert

• added_size – Optional size of insertion (calculated if None)

Returns:

Updated heap instance

key_buffer: Array | ndarray | bool | number

key_store: Array | ndarray | bool | number

static make_batched(key: Array | ndarray | bool | number, val: Xtructurable, batch_size: int)

Convert unbatched arrays into batched format suitable for the queue.

Parameters:

• key – Array of keys to batch

• val – Xtructurable of values to batch

• batch_size – Desired batch size

Returns:

• Batched key array

• Batched value array

Return type:

tuple containing

max_size: int

merge_buffer(blockk: Array | ndarray | bool | number, blockv: Xtructurable)

Merge buffer contents with block contents, handling overflow conditions.

This method is crucial for maintaining the heap property when inserting new elements. It handles the case where the buffer might overflow into the main storage.

Parameters:

• blockk – Block keys array

• blockv – Block values

• bufferk – Buffer keys array

• bufferv – Buffer values

Returns:

• Updated block keys

• Updated block values

• Updated buffer keys

• Updated buffer values

• Boolean indicating if buffer overflow occurred

Return type:

tuple containing

replace(**kwargs)

property size

to_tuple()

val_buffer: Xtructurable

val_store: Xtructurable