xtructure.hashtable package

Submodules

xtructure.hashtable.constants module

xtructure.hashtable.hash_utils module

Hash helpers for bucketed double hashing.

xtructure.hashtable.hash_utils.get_new_idx_byterized(input: Xtructurable, modulus: int, seed: int) → tuple[Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number]

Hash a Xtructurable and return index, step, uint32ed, fingerprint, and hash pair.

xtructure.hashtable.hash_utils.get_new_idx_from_uint32ed(input_uint32ed: Array | ndarray | bool | number, modulus: int, seed: int) → tuple[Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number]

Calculate a new hash bucket index, probe step, and both hash values from a uint32ed.

xtructure.hashtable.hash_utils.get_new_idx_hashed(input: Xtructurable, modulus: int, seed: int) → tuple[Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number]

Hash a Xtructurable and return index, step, fingerprint, and hash pair.

This avoids materializing/returning the (potentially very wide) uint32ed buffer.

xtructure.hashtable.insert module

Insertion helpers for HashTable.

xtructure.hashtable.insert_pallas module

xtructure.hashtable.insert_pallas.pallas_insert_enabled() → bool

xtructure.hashtable.insert_pallas.reserve_slots_pallas(bucket_fill_levels: Array | ndarray | bool | number, bucket_occupancy: Array | ndarray | bool | number, start_buckets: Array | ndarray | bool | number, probe_steps: Array | ndarray | bool | number, active: Array | ndarray | bool | number, *, bucket_size: int, capacity: int) → Tuple[Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number]

xtructure.hashtable.insert_triton module

xtructure.hashtable.insert_triton.reserve_slots_triton(bucket_occupancy: Array | ndarray | bool | number, start_buckets: Array | ndarray | bool | number, probe_steps: Array | ndarray | bool | number, active: Array | ndarray | bool | number, *, bucket_size: int, capacity: int) → Tuple[Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number]

xtructure.hashtable.insert_triton.triton_insert_enabled() → bool

xtructure.hashtable.lookup module

Lookup helpers for HashTable.

xtructure.hashtable.table module

HashTable data container and public API.

class xtructure.hashtable.table.HashTable(seed: int, capacity: int, _capacity: int, bucket_size: int, size: int, table: Xtructurable, bucket_fill_levels: Array | ndarray | bool | number, bucket_occupancy: Array | ndarray | bool | number, fingerprints: Array | ndarray | bool | number, max_probes: int)

Bases: object

Bucketed Double Hash Table Implementation

Uses double hashing with buckets to resolve collisions.

bucket_fill_levels: Array | ndarray | bool | number

bucket_occupancy: Array | ndarray | bool | number

bucket_size: int

static build(dataclass: Xtructurable, seed: int, capacity: int, bucket_size: int = 8, hash_size_multiplier: int = 2, max_probes: int | None = None) → HashTable

Initialize a new hash table backed by JAX-friendly storage.

capacity: int

fingerprints: Array | ndarray | bool | number

from_tuple()

insert(input: Xtructurable) → tuple[HashTable, bool, Xtructurable]

lookup(input: Xtructurable) → tuple[Xtructurable, bool]

lookup_bucket(input: Xtructurable) → tuple[Xtructurable, Array | ndarray | bool | number, Array | ndarray | bool | number]

lookup_parallel(inputs: Xtructurable, filled: Array | ndarray | bool | number | bool = True) → tuple[Xtructurable, Array | ndarray | bool | number]

max_probes: int

parallel_insert(inputs: Xtructurable, filled: Array | ndarray | bool | number | bool | None = None, unique_key: Array | ndarray | bool | number | None = None)

replace(**kwargs)

seed: int

size: int

table: Xtructurable

to_tuple()

xtructure.hashtable.types module

class xtructure.hashtable.types.BucketIdx(index: Annotated[Union[jax.jaxlib._jax.Array, numpy.ndarray, numpy.bool, numpy.number], FieldDescriptor(dtype=<class 'jax.numpy.uint32'>, fill_value=4294967295, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)], slot_index: Annotated[Union[jax.jaxlib._jax.Array, numpy.ndarray, numpy.bool, numpy.number], FieldDescriptor(dtype=<class 'jax.numpy.uint8'>, fill_value=255, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)])

Bases: object

allclose(b: Any, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) → bool | Array

Returns True if two arrays are element-wise equal within a tolerance.

astype(dtype: Any, copy: bool = False, device: Any = None) → T

Copy of the array, cast to a specified type.

property at

property batch_shape

block() → Any

Assemble an nd-array from nested lists of blocks.

broadcast_to(shape: Sequence[int]) → T

Broadcast an array to a new shape.

property bytes

Convert entire state tree to flattened byte array.

check_invariants()

column_stack() → Any

Stack 1-D arrays as columns into a 2-D array.

classmethod default(shape: Tuple[int, ...] = ()) → T

default_dtype = (<class 'jax.numpy.uint32'>, <class 'jax.numpy.uint8'>)

default_shape = ((), ())

dstack(dtype: Any = None) → Any

Stack arrays in sequence depth wise (along third axis).

property dtype: dtype

Get dtypes of all fields in the dataclass

equal(y: Any) → T

Return (x == y) element-wise.

expand_dims(axis: int) → T

Insert a new axis into every field.

flatten() → T

Flatten the batch dimensions of a dataclass instance.

flip(axis: int | Sequence[int] | None = None) → T

Reverse the order of elements in an array along the given axis.

from_tuple()

hash(seed=0)

Main hash function that converts state to uint32 lanes and hashes them.

hash_pair(seed=0)

Hash function that returns two 32-bit hashes.

hash_pair_with_uint32ed(seed=0)

Hash function that returns two 32-bit hashes and the uint32 lanes.

hash_with_uint32ed(seed=0)

Main hash function that converts state to uint32 lanes and hashes them. Returns both hash value and its uint32 representation.

hstack(dtype: Any = None) → Any

Stack arrays in sequence horizontally (column wise).

index: uint32'>, fill_value=4294967295, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)]

is_xtructed = True

isclose(b: Any, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) → T

Returns a boolean array where two arrays are element-wise equal within a tolerance.

classmethod load(path: str) → T

Loads an instance from a .npz file.

moveaxis(source: int | Sequence[int], destination: int | Sequence[int]) → T

Move axes of an array to new positions.

property ndim: int

Return number of batch dimensions for structured instances.

not_equal(y: Any) → T

Return (x != y) element-wise.

pad(pad_width: int | tuple[int, ...] | tuple[tuple[int, int], ...], mode: str = 'constant', **kwargs) → T

Pad xtructure dataclasses using a jnp.pad compatible interface.

classmethod random(shape=(), key=None)

replace(**kwargs)

reshape(new_shape: tuple[int, ...] | int, *args: int) → T

Reshape the batch dimensions of a dataclass instance.

Supports both reshape(instance, (2, 3)) and reshape(instance, 2, 3) syntax. Also supports -1 for dimension inference.

roll(shift: int | Sequence[int], axis: int | Sequence[int] | None = None) → T

Roll array elements along a given axis.

rot90(k: int = 1, axes: tuple[int, int] = (0, 1)) → T

Rotate an array by 90 degrees in the plane specified by axes.

save(path: str, *, packed: bool = True)

Saves the instance to a .npz file.

property shape: shape

Returns a namedtuple containing the batch shape (if present) and the shapes of all fields. If a field is itself a xtructure_dataclass, its shape is included as a nested namedtuple.

slot_index: uint8'>, fill_value=255, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)]

squeeze(axis: int | tuple[int, ...] | None = None) → T

Remove axes of length one from every field.

str(**kwargs)

property structured_type: StructuredType

swapaxes(axis1: int, axis2: int) → T

Swap two batch axes.

to_tuple()

transpose(axes: tuple[int, ...] | None = None) → T

Transpose batch dimensions of every field.

property uint32ed

Convert pytree to uint32 array.

vstack(dtype: Any = None) → Any

Stack arrays in sequence vertically (row wise).

class xtructure.hashtable.types.HashIdx(index: Annotated[Union[jax.jaxlib._jax.Array, numpy.ndarray, numpy.bool, numpy.number], FieldDescriptor(dtype=<class 'jax.numpy.uint32'>, fill_value=4294967295, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)])

Bases: object

allclose(b: Any, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) → bool | Array

Returns True if two arrays are element-wise equal within a tolerance.

astype(dtype: Any, copy: bool = False, device: Any = None) → T

Copy of the array, cast to a specified type.

property at

property batch_shape

block() → Any

Assemble an nd-array from nested lists of blocks.

broadcast_to(shape: Sequence[int]) → T

Broadcast an array to a new shape.

property bytes

Convert entire state tree to flattened byte array.

check_invariants()

column_stack() → Any

Stack 1-D arrays as columns into a 2-D array.

classmethod default(shape: Tuple[int, ...] = ()) → T

default_dtype = (<class 'jax.numpy.uint32'>,)

default_shape = ((),)

dstack(dtype: Any = None) → Any

Stack arrays in sequence depth wise (along third axis).

property dtype: dtype

Get dtypes of all fields in the dataclass

equal(y: Any) → T

Return (x == y) element-wise.

expand_dims(axis: int) → T

Insert a new axis into every field.

flatten() → T

Flatten the batch dimensions of a dataclass instance.

flip(axis: int | Sequence[int] | None = None) → T

Reverse the order of elements in an array along the given axis.

from_tuple()

hash(seed=0)

Main hash function that converts state to uint32 lanes and hashes them.

hash_pair(seed=0)

Hash function that returns two 32-bit hashes.

hash_pair_with_uint32ed(seed=0)

Hash function that returns two 32-bit hashes and the uint32 lanes.

hash_with_uint32ed(seed=0)

Main hash function that converts state to uint32 lanes and hashes them. Returns both hash value and its uint32 representation.

hstack(dtype: Any = None) → Any

Stack arrays in sequence horizontally (column wise).

index: uint32'>, fill_value=4294967295, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)]

is_xtructed = True

isclose(b: Any, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) → T

Returns a boolean array where two arrays are element-wise equal within a tolerance.

classmethod load(path: str) → T

Loads an instance from a .npz file.

moveaxis(source: int | Sequence[int], destination: int | Sequence[int]) → T

Move axes of an array to new positions.

property ndim: int

Return number of batch dimensions for structured instances.

not_equal(y: Any) → T

Return (x != y) element-wise.

pad(pad_width: int | tuple[int, ...] | tuple[tuple[int, int], ...], mode: str = 'constant', **kwargs) → T

Pad xtructure dataclasses using a jnp.pad compatible interface.

classmethod random(shape=(), key=None)

replace(**kwargs)

reshape(new_shape: tuple[int, ...] | int, *args: int) → T

Reshape the batch dimensions of a dataclass instance.

Supports both reshape(instance, (2, 3)) and reshape(instance, 2, 3) syntax. Also supports -1 for dimension inference.

roll(shift: int | Sequence[int], axis: int | Sequence[int] | None = None) → T

Roll array elements along a given axis.

rot90(k: int = 1, axes: tuple[int, int] = (0, 1)) → T

Rotate an array by 90 degrees in the plane specified by axes.

save(path: str, *, packed: bool = True)

Saves the instance to a .npz file.

property shape: shape

Returns a namedtuple containing the batch shape (if present) and the shapes of all fields. If a field is itself a xtructure_dataclass, its shape is included as a nested namedtuple.

squeeze(axis: int | tuple[int, ...] | None = None) → T

Remove axes of length one from every field.

str(**kwargs)

property structured_type: StructuredType

swapaxes(axis1: int, axis2: int) → T

Swap two batch axes.

to_tuple()

transpose(axes: tuple[int, ...] | None = None) → T

Transpose batch dimensions of every field.

property uint32ed

Convert pytree to uint32 array.

vstack(dtype: Any = None) → Any

Stack arrays in sequence vertically (row wise).

Module contents

class xtructure.hashtable.HashIdx(index: Annotated[Union[jax.jaxlib._jax.Array, numpy.ndarray, numpy.bool, numpy.number], FieldDescriptor(dtype=<class 'jax.numpy.uint32'>, fill_value=4294967295, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)])

Bases: object

allclose(b: Any, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) → bool | Array

Returns True if two arrays are element-wise equal within a tolerance.

astype(dtype: Any, copy: bool = False, device: Any = None) → T

Copy of the array, cast to a specified type.

property at

property batch_shape

block() → Any

Assemble an nd-array from nested lists of blocks.

broadcast_to(shape: Sequence[int]) → T

Broadcast an array to a new shape.

property bytes

Convert entire state tree to flattened byte array.

check_invariants()

column_stack() → Any

Stack 1-D arrays as columns into a 2-D array.

classmethod default(shape: Tuple[int, ...] = ()) → T

default_dtype = (<class 'jax.numpy.uint32'>,)

default_shape = ((),)

dstack(dtype: Any = None) → Any

Stack arrays in sequence depth wise (along third axis).

property dtype: dtype

Get dtypes of all fields in the dataclass

equal(y: Any) → T

Return (x == y) element-wise.

expand_dims(axis: int) → T

Insert a new axis into every field.

flatten() → T

Flatten the batch dimensions of a dataclass instance.

flip(axis: int | Sequence[int] | None = None) → T

Reverse the order of elements in an array along the given axis.

from_tuple()

hash(seed=0)

Main hash function that converts state to uint32 lanes and hashes them.

hash_pair(seed=0)

Hash function that returns two 32-bit hashes.

hash_pair_with_uint32ed(seed=0)

Hash function that returns two 32-bit hashes and the uint32 lanes.

hash_with_uint32ed(seed=0)

Main hash function that converts state to uint32 lanes and hashes them. Returns both hash value and its uint32 representation.

hstack(dtype: Any = None) → Any

Stack arrays in sequence horizontally (column wise).

index: uint32'>, fill_value=4294967295, intrinsic_shape=(), bits=None, packed_bits=None, unpacked_dtype=None, unpacked_intrinsic_shape=None, fill_value_factory=None, validator=None)]

is_xtructed = True

isclose(b: Any, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) → T

Returns a boolean array where two arrays are element-wise equal within a tolerance.

classmethod load(path: str) → T

Loads an instance from a .npz file.

moveaxis(source: int | Sequence[int], destination: int | Sequence[int]) → T

Move axes of an array to new positions.

property ndim: int

Return number of batch dimensions for structured instances.

not_equal(y: Any) → T

Return (x != y) element-wise.

pad(pad_width: int | tuple[int, ...] | tuple[tuple[int, int], ...], mode: str = 'constant', **kwargs) → T

Pad xtructure dataclasses using a jnp.pad compatible interface.

classmethod random(shape=(), key=None)

replace(**kwargs)

reshape(new_shape: tuple[int, ...] | int, *args: int) → T

Reshape the batch dimensions of a dataclass instance.

Supports both reshape(instance, (2, 3)) and reshape(instance, 2, 3) syntax. Also supports -1 for dimension inference.

roll(shift: int | Sequence[int], axis: int | Sequence[int] | None = None) → T

Roll array elements along a given axis.

rot90(k: int = 1, axes: tuple[int, int] = (0, 1)) → T

Rotate an array by 90 degrees in the plane specified by axes.

save(path: str, *, packed: bool = True)

Saves the instance to a .npz file.

property shape: shape

Returns a namedtuple containing the batch shape (if present) and the shapes of all fields. If a field is itself a xtructure_dataclass, its shape is included as a nested namedtuple.

squeeze(axis: int | tuple[int, ...] | None = None) → T

Remove axes of length one from every field.

str(**kwargs)

property structured_type: StructuredType

swapaxes(axis1: int, axis2: int) → T

Swap two batch axes.

to_tuple()

transpose(axes: tuple[int, ...] | None = None) → T

Transpose batch dimensions of every field.

property uint32ed

Convert pytree to uint32 array.

vstack(dtype: Any = None) → Any

Stack arrays in sequence vertically (row wise).

class xtructure.hashtable.HashTable(seed: int, capacity: int, _capacity: int, bucket_size: int, size: int, table: Xtructurable, bucket_fill_levels: Array | ndarray | bool | number, bucket_occupancy: Array | ndarray | bool | number, fingerprints: Array | ndarray | bool | number, max_probes: int)

Bases: object

Bucketed Double Hash Table Implementation

Uses double hashing with buckets to resolve collisions.

bucket_fill_levels: Array | ndarray | bool | number

bucket_occupancy: Array | ndarray | bool | number

bucket_size: int

static build(dataclass: Xtructurable, seed: int, capacity: int, bucket_size: int = 8, hash_size_multiplier: int = 2, max_probes: int | None = None) → HashTable

Initialize a new hash table backed by JAX-friendly storage.

capacity: int

fingerprints: Array | ndarray | bool | number

from_tuple()

insert(input: Xtructurable) → tuple[HashTable, bool, Xtructurable]

lookup(input: Xtructurable) → tuple[Xtructurable, bool]

lookup_bucket(input: Xtructurable) → tuple[Xtructurable, Array | ndarray | bool | number, Array | ndarray | bool | number]

lookup_parallel(inputs: Xtructurable, filled: Array | ndarray | bool | number | bool = True) → tuple[Xtructurable, Array | ndarray | bool | number]

max_probes: int

parallel_insert(inputs: Xtructurable, filled: Array | ndarray | bool | number | bool | None = None, unique_key: Array | ndarray | bool | number | None = None)

replace(**kwargs)

seed: int

size: int

table: Xtructurable

to_tuple()