archimedes.tree.struct

Utilities for defining custom classes that can be used with pytree transformations.

This module provides tools for creating structured data types that work seamlessly with Archimedes’ pytree functions. These tools are built on Python’s dataclasses with extensions for pytree-specific behavior.

The module re-exports several names from the dataclasses module:

InitVarType annotation for init-only variables in dataclasses

Used to mark fields that should be passed to __post_init__ but not stored.

fieldsFunction to retrieve fields of a dataclass

Returns a list of Field objects representing the fields of the dataclass. This is useful for introspection and validation of dataclass instances.

replaceFunction to create a new dataclass instance with updated fields

For pytree nodes created with @pytree_node, use the .replace() method instead.

Functions

dataclass_transform(*[, eq_default, ...])	Decorator to mark an object as providing dataclass-like behaviour.
field([static, metadata])	Create a field specification with pytree-related metadata.
fields(class_or_instance)	Return a tuple describing the fields of this dataclass.
is_pytree_node(obj)	Check if an object is a registered pytree node.
pytree_node([cls])	Decorator to convert a class into a frozen dataclass registered as a pytree node.
register_dataclass(nodetype[, data_fields, ...])	Register a dataclass as a pytree node with customized field handling.
replace(obj, /, **changes)	Return a new object replacing specified fields with new values.

Classes

Any(*args, **kwargs)	Special type indicating an unconstrained type.
Callable()
InitVar(type)
TypeVar	Type variable.