xtructure.core.xtructure_decorators.aggregate_bitpack package

Submodules

xtructure.core.xtructure_decorators.aggregate_bitpack.bitpack module

Bitpacking utilities for compact serialization.

This module packs arrays whose values use only a small number of active bits into a compact uint8 stream, and unpacks them back to arrays.

Design notes: - Supports 1..8 bits per value (inclusive). - Uses JAX-compatible primitives (jnp/jax.vmap) so it can run on device.

xtructure.core.xtructure_decorators.aggregate_bitpack.bitpack.from_uint8(packed_bytes: Array | ndarray | bool | number, target_shape: tuple[int, ...], active_bits: int = 1) → Array | ndarray | bool | number[source]

Unpack a uint8 stream back into an array of shape target_shape.

Notes:

  • For active_bits==1, returns bool.

  • For active_bits>1, returns uint8 values in [0, 2**active_bits - 1]. Caller can cast to a desired integer dtype.

xtructure.core.xtructure_decorators.aggregate_bitpack.bitpack.packed_num_bytes(num_values: int, active_bits: int) → int[source]

Return the number of uint8 bytes required to pack num_values values.

This mirrors the packing strategy of to_uint8() (block-aligned for 3/5/6/7 bits).

xtructure.core.xtructure_decorators.aggregate_bitpack.bitpack.to_uint8(values: Array | ndarray | bool | number, active_bits: int = 1) → Array | ndarray | bool | number[source]

Pack an array into a uint8 stream using active_bits per value.

Parameters:

  • values – Input array. For active_bits==1, can be bool or integer (0/!=0). For active_bits>1, must be integer.

  • active_bits – Bits per value in [1, 8].

Returns:

A 1D uint8 array of packed bytes.

xtructure.core.xtructure_decorators.aggregate_bitpack.kernels module

Aggregate bitpack kernels.

This module contains a vectorized XLA implementation and a pallas-first implementation for packing the aggregate bitstream.

The packing kernel computes each output word independently from a precomputed contribution table, avoiding per-value scatter updates.

xtructure.core.xtructure_decorators.aggregate_bitpack.kernels.pack_words_all_pallas(values_stream_u32: Array, tables: _AggWordContribTables, *, backend: Literal['triton', 'mosaic_tpu'], word_tile: int) → Array[source]
xtructure.core.xtructure_decorators.aggregate_bitpack.kernels.pack_words_all_xla(values_stream_u32: Array, tables: _AggWordContribTables) → Array[source]

Pack (flat_n, total_values) uint32 stream -> (flat_n, words_all_len) uint32 words.

xtructure.core.xtructure_decorators.aggregate_bitpack.spec module

Aggregate bitpack layout and specification helpers.

xtructure.core.xtructure_decorators.aggregate_bitpack.view module

Logical unpacked view class builder for aggregate bitpack.

xtructure.core.xtructure_decorators.aggregate_bitpack.view.build_unpacked_view_cls(root_cls: type, *, default_unpack_dtype: Callable[[int], Any]) → tuple[type, dict[type, type]][source]

Return the logical view type for the aggregate-packed class and a cache.

Module contents

Aggregate bitpacking across fields of a dataclass.

This decorator adds a .packed property that returns a packed representation containing a word-aligned uint32 stream plus an optional uint8 tail, and a .unpacked property on the packed representation that reconstructs a logical view.

Opt-in via @xtructure_dataclass(aggregate_bitpack=True).

Rules: - Only pack primitive array-like fields whose FieldDescriptor.bits is set. - bits can be 1..32. - Nested xtructure_dataclass fields are supported for scalar nested fields (intrinsic_shape == ()).

xtructure.core.xtructure_decorators.aggregate_bitpack.add_aggregate_bitpack(cls: Type[T]) → Type[T][source]