xtructure.core.xtructure_decorators package

Subpackages

• xtructure.core.xtructure_decorators.aggregate_bitpack package
  • Submodules
  • xtructure.core.xtructure_decorators.aggregate_bitpack.bitpack module
    • from_uint8()
    • packed_num_bytes()
    • to_uint8()
  • xtructure.core.xtructure_decorators.aggregate_bitpack.kernels module
    • pack_words_all_pallas()
    • pack_words_all_xla()
  • xtructure.core.xtructure_decorators.aggregate_bitpack.spec module
  • xtructure.core.xtructure_decorators.aggregate_bitpack.view module
    • build_unpacked_view_cls()
  • Module contents
    • add_aggregate_bitpack()

Submodules

xtructure.core.xtructure_decorators.annotate module

xtructure.core.xtructure_decorators.bitpack_accessors module

Decorator that adds in-memory bitpack accessors for packed fields.

Packed fields are those whose FieldDescriptor declares packed_bits plus unpacked_dtype and unpacked_intrinsic_shape. The stored leaf is expected to be a uint8 byte-stream (typically created via FieldDescriptor.packed_tensor), and the decorator adds:

• <field>_unpacked property: returns logical array of shape batch + unpacked_shape

• set_unpacked(**kwargs) method: packs provided logical arrays and returns a new instance

xtructure.core.xtructure_decorators.bitpack_accessors.add_bitpack_accessors(cls: Type[T]) → Type[T]

Attach unpacked accessors for any packed fields defined on the class.

xtructure.core.xtructure_decorators.default module

class xtructure.core.xtructure_decorators.default.FieldInfo(name: str, field_type: str, descriptor: FieldDescriptor | None, dtype: Any, fill_value: Any, fill_value_factory: Callable[[Tuple[int, ...], Any], Any] | None, intrinsic_shape: Tuple[int, ...], nested_class_type: Type | None)

Bases: NamedTuple

Pre-computed field information for efficient default value generation.

descriptor: FieldDescriptor | None

Alias for field number 2

dtype: Any

Alias for field number 3

field_type: str

Alias for field number 1

fill_value: Any

Alias for field number 4

fill_value_factory: Callable[[Tuple[int, ...], Any], Any] | None

Alias for field number 5

intrinsic_shape: Tuple[int, ...]

Alias for field number 6

name: str

Alias for field number 0

nested_class_type: Type | None

Alias for field number 7

xtructure.core.xtructure_decorators.default.add_default_method(cls: Type[T]) → Type[T]

xtructure.core.xtructure_decorators.hash module

xtructure.core.xtructure_decorators.hash.byterize_hash_func_builder(x: Xtructurable)

Build a hash function for the pytree. This function creates a JIT-compiled hash function that converts pytree leaves to uint32 lanes and then reduces them with a vectorized avalanche hash.

Parameters:

x – Example pytree to determine the structure

Returns:

JIT-compiled hash function that takes a pytree and seed

xtructure.core.xtructure_decorators.hash.hash_fast_uint32ed(uint32ed, seed=Array(0, dtype=uint32))

Vectorized hash reducer for uint32 streams.

xtructure.core.xtructure_decorators.hash.hash_fast_uint32ed_pair(uint32ed, seed=Array(0, dtype=uint32))

Vectorized hash reducer for uint32 streams (returns two 32-bit hashes).

This is intended to support double-hashing and wide signatures without requiring a second full pass over the input.

xtructure.core.xtructure_decorators.hash.hash_function_decorator(cls)

Decorator to add a hash function to a class.

xtructure.core.xtructure_decorators.hash.uint32ed_to_hash(uint32ed, seed)

Convert uint32 array to hash value.

xtructure.core.xtructure_decorators.hash.uint32ed_to_hash_pair(uint32ed, seed)

Convert uint32 array to a pair of 32-bit hashes.

xtructure.core.xtructure_decorators.indexing module

class xtructure.core.xtructure_decorators.indexing.AtIndexer(obj_instance)

Bases: object

xtructure.core.xtructure_decorators.indexing.add_indexing_methods(cls: Type[T]) → Type[T]

Augments the class with an __getitem__ method for indexing/slicing and an at property that enables JAX-like out-of-place updates (e.g., instance.at[index].set(value)).

The __getitem__ method allows instances to be indexed, applying the index to each field. The at property provides access to an updater object for specific indices.

xtructure.core.xtructure_decorators.io module

Decorator to add save and load methods to xtructure dataclasses.

xtructure.core.xtructure_decorators.io.add_io_methods(cls: Type[T]) → Type[T]

Augments the class with save and load methods.

The save method allows an instance to be saved to a file. The load method allows the class to load an instance from a file.

xtructure.core.xtructure_decorators.method_factory module

Instance method factory for xtructure dataclasses.

This module provides utilities to automatically generate instance methods from xnp (xtructure_numpy) functions. When a new xnp function is added, simply add it to the appropriate list below to make it available as an instance method on all xtructure dataclasses.

xtructure.core.xtructure_decorators.method_factory.add_xnp_instance_methods(cls: Type[T]) → Type[T]

Add instance methods that delegate to xnp functions.

This decorator adds methods like .roll(), .flip(), .astype() etc. that call the corresponding xnp.roll(), xnp.flip(), xnp.astype() functions.

To add a new method when a new xnp function is created: 1. Add the function name and import to _XNP_METHOD_REGISTRY below 2. The method will automatically be available on all xtructure dataclasses

xtructure.core.xtructure_decorators.ops module

xtructure.core.xtructure_decorators.ops.add_comparison_operators(cls: Type[T]) → Type[T]

Adds custom __eq__ and __ne__ methods to the class.

Semantics:

• For two instances of the same xtructure dataclass type, comparisons are performed field-wise and then reduced to a single boolean:

  • __eq__: True iff all fields are equal (via jnp.all(x == y) per field).

  • __ne__: True iff any field is different (via jnp.any(x != y) per field).

• If other is not the same type, returns NotImplemented.

Note: these operators return a scalar boolean (JAX bool array / Python bool), not a dataclass of booleans.

xtructure.core.xtructure_decorators.shape module

xtructure.core.xtructure_decorators.shape.add_shape_dtype_len(cls: Type[T]) → Type[T]

Augments the class with shape and dtype properties to inspect its fields, and a __len__ method.

The shape and dtype properties return namedtuples reflecting the structure of the dataclass fields. The __len__ method conventionally returns the size of the first dimension of the first field of the instance, which is often useful for determining batch sizes.

xtructure.core.xtructure_decorators.string_format module

xtructure.core.xtructure_decorators.string_format.add_string_representation_methods(cls: Type[T]) → Type[T]

Adds custom __str__ and str methods to the class for generating a more informative string representation.

It handles instances categorized by structured_type differently:

• SINGLE: Uses the original __str__ (or repr if basic) of the instance.

• BATCHED: Provides a summarized view if the batch is large, showing the first few and last few elements, along with the batch shape. Uses rich for formatting.

• UNSTRUCTURED: Indicates that the data is unstructured relative to its default shape.

xtructure.core.xtructure_decorators.structure_util module

xtructure.core.xtructure_decorators.structure_util.add_structure_utilities(cls: Type[T]) → Type[T]

Augments the class with utility methods and properties related to its structural representation (based on a ‘default’ instance), batch operations, and random instance generation.

Requires the class to have a default classmethod, which is used to determine default shapes, dtypes, and behaviors.

Adds:

• Properties:

  • default_shape: Shape of the instance returned by cls.default().

  • structured_type: An enum (StructuredType) indicating if the instance is SINGLE, BATCHED, or UNSTRUCTURED relative to its default shape.

  • batch_shape: The shape of the batch dimensions if structured_type is BATCHED.

• Instance Methods:

  • reshape(*new_shape): Reshapes the batch dimensions of a BATCHED instance.

  • flatten(): Flattens the batch dimensions of a BATCHED instance.

  • transpose(axes=None): Transposes only the batch dimensions.

• Classmethod:

  • random(shape=(), key=None): Generates an instance with random data. The shape argument specifies the desired batch shape, which is prepended to the default field shapes.

xtructure.core.xtructure_decorators.validation module

xtructure.core.xtructure_decorators.validation.add_runtime_validation(cls: Type[T], *, enabled: bool) → Type[T]

Inject a post-init validator that checks dtype and trailing shape.

Module contents

xtructure.core.xtructure_decorators.xtructure_dataclass(cls: Type[T] | None = None, *, validate: bool = False, aggregate_bitpack: bool = False, bitpack: str = 'auto') → Callable[[Type[T]], Type[Xtructurable[T]]] | Type[Xtructurable[T]]

Decorator that ensures the input class is a base_dataclass (or converts it to one) and then augments it with additional functionality related to its structure, type, and operations like indexing, default instance creation, random instance generation, and string representation.

It adds properties like shape, dtype, default_shape, structured_type, batch_shape, and methods like __getitem__, __len__, reshape, flatten, transpose, random, and __str__.

Parameters:

• cls – The class to be decorated. It is expected to have a default classmethod for some functionalities.

• validate – When True, injects a runtime validator that checks field dtypes and trailing shapes after every instantiation.

Returns:

The decorated class with the aforementioned additional functionalities.