xtructure.core.xtructure_decorators package

Submodules

xtructure.core.xtructure_decorators.annotate module

xtructure.core.xtructure_decorators.default module

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_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.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.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.aggregate_bitpack module

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]

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.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]

Attach shape/dtype/structured-type properties using Xtructure Layout facts.

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.