xtructure.core.xtructure_decorators package

Submodules

xtructure.core.xtructure_decorators.annotate module

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.default.is_xtructure_class(dtype: Any) → bool

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.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. These methods perform element-wise comparisons on the fields of the dataclass and return a new instance of the class containing boolean arrays.

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.

• 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.structure_util.is_nested_xtructure(dtype: Any) → bool

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) → 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, 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.