xtructure.bgpq package

Subpackages

• xtructure.bgpq.merge_split package
  • Submodules
  • xtructure.bgpq.merge_split.benchmark_merge_split module
    • BenchValue
      • BenchValue.a
      • BenchValue.allclose()
      • BenchValue.astype()
      • BenchValue.at
      • BenchValue.b
      • BenchValue.batch_shape
      • BenchValue.block()
      • BenchValue.broadcast_to()
      • BenchValue.bytes
      • BenchValue.c
      • BenchValue.check_invariants()
      • BenchValue.column_stack()
      • BenchValue.default()
      • BenchValue.default_dtype
      • BenchValue.default_shape
      • BenchValue.dstack()
      • BenchValue.dtype
      • BenchValue.equal()
      • BenchValue.expand_dims()
      • BenchValue.flatten()
      • BenchValue.flip()
      • BenchValue.from_tuple()
      • BenchValue.hash()
      • BenchValue.hash_pair()
      • BenchValue.hash_pair_with_uint32ed()
      • BenchValue.hash_with_uint32ed()
      • BenchValue.hstack()
      • BenchValue.is_xtructed
      • BenchValue.isclose()
      • BenchValue.load()
      • BenchValue.moveaxis()
      • BenchValue.ndim
      • BenchValue.not_equal()
      • BenchValue.pad()
      • BenchValue.random()
      • BenchValue.replace()
      • BenchValue.reshape()
      • BenchValue.roll()
      • BenchValue.rot90()
      • BenchValue.save()
      • BenchValue.shape
      • BenchValue.squeeze()
      • BenchValue.str()
      • BenchValue.structured_type
      • BenchValue.swapaxes()
      • BenchValue.to_tuple()
      • BenchValue.transpose()
      • BenchValue.uint32ed
      • BenchValue.vstack()
    • main()
    • run_bench()
    • run_bench_values()
  • 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_arrays_parallel_kv()
  • 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.jaxlib._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.jaxlib._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'>) → BGPQ

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

delete_mins()

Remove and return the minimum elements from the queue.

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) → BGPQ

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

Parameters:

• 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

make_batched_like(key: Array | ndarray | bool | number, val: Xtructurable)

Pad key/val to this heap’s batch_size (a static_fields config).

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

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

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.jaxlib._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.jaxlib._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'>) → BGPQ

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

delete_mins()

Remove and return the minimum elements from the queue.

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) → BGPQ

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

Parameters:

• 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

make_batched_like(key: Array | ndarray | bool | number, val: Xtructurable)

Pad key/val to this heap’s batch_size (a static_fields config).

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

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