From 21701420123672a2c862d8ac51ed94b81ee774ea Mon Sep 17 00:00:00 2001 From: renaissancenerd Date: Thu, 16 Apr 2026 13:46:21 +0200 Subject: [PATCH] updating dataclasses.replace and register_jax_tree --- src/drinx/base.py | 63 +++++++++++++++++++++++++++++++++++++++--- src/drinx/transform.py | 45 ++++++++++++++++++++++++++---- src/drinx/visualize.py | 4 +-- 3 files changed, 99 insertions(+), 13 deletions(-) diff --git a/src/drinx/base.py b/src/drinx/base.py index 912c29a..d738b02 100644 --- a/src/drinx/base.py +++ b/src/drinx/base.py @@ -16,6 +16,60 @@ ) +def _dataclass_replace(obj: Any, **changes: Any) -> Any: + """Functional replacement for :func:`dataclasses.replace` that correctly + handles private (``init=False``) fields. + + :func:`dataclasses.replace` only forwards ``init=True`` fields to + ``__init__``, silently dropping any ``init=False`` field that appears in + *changes*. This wrapper: + + 1. Collects the current values of **all** fields on *obj*. + 2. Applies *changes* on top (regardless of ``init`` status). + 3. Passes only the ``init=True`` fields to ``__init__`` to construct the + new instance. + 4. Uses ``object.__setattr__`` to stamp the ``init=False`` fields onto the + freshly created (still-unfrozen) object. + + Args: + obj: The frozen dataclass instance to copy. + **changes: Field-name → new-value pairs. Both ``init=True`` and + ``init=False`` fields are accepted. + + Returns: + A new instance of ``type(obj)`` with the requested fields replaced. + + Raises: + TypeError: If any key in *changes* is not a recognised field name. + """ + all_fields = dataclasses.fields(obj) + known_names = {f.name for f in all_fields} + unknown = set(changes) - known_names + if unknown: + raise TypeError(f"_dataclass_replace() got unexpected field names: {unknown!r}") + + # Collect current values for every field, then overlay changes + current_values: dict[str, Any] = { + f.name: object.__getattribute__(obj, f.name) for f in all_fields + } + current_values.update(changes) + + init_kwargs = {f.name: current_values[f.name] for f in all_fields if f.init} + non_init_overrides = { + f.name: current_values[f.name] + for f in all_fields + if not f.init and f.name in changes + } + + new_obj = type(obj)(**init_kwargs) + + # Stamp non-init fields that were explicitly changed + for name, value in non_init_overrides.items(): + object.__setattr__(new_obj, name, value) + + return new_obj + + @dataclass_transform( field_specifiers=( orig_field, @@ -438,8 +492,9 @@ def aset( f"Can only set attribute functionally on a dataclass, but got {current_parent.__class__}" ) - # Use standard dataclasses.replace to functionally copy and update the frozen dataclass - cur_attr = dataclasses.replace(current_parent, **{str(op): cur_attr}) + # Use _dataclass_replace (instead of dataclasses.replace) so that + # private/non-init fields are handled correctly. + cur_attr = _dataclass_replace(current_parent, **{str(op): cur_attr}) elif op_type in ("index", "key"): if not hasattr(current_parent, "copy"): @@ -520,8 +575,8 @@ def updated_copy(self, **kwargs: Any) -> Self: Returns: Self: A newly instantiated object with the updated attributes. """ - # Directly utilize dataclasses.replace for standard functional updates - return dataclasses.replace(self, **kwargs) + # Use _dataclass_replace so private (init=False) fields are handled correctly. + return _dataclass_replace(self, **kwargs) _DC = TypeVar("_DC", bound="DataClass") diff --git a/src/drinx/transform.py b/src/drinx/transform.py index 0f22425..72f7eb1 100644 --- a/src/drinx/transform.py +++ b/src/drinx/transform.py @@ -10,24 +10,57 @@ def _register_jax_tree(cls_: type[T]) -> type[T]: - """Registers a class as a JAX Pytree, safely preventing double-registration.""" + """Registers a class as a JAX Pytree, safely preventing double-registration. + + Field classification: + - ``jax_static=True`` + ``init=True`` -> aux bucket 1 (restored via __init__) + - ``jax_static=False`` + ``init=True`` -> traced leaves (restored via __init__) + - any + ``init=False`` -> aux bucket 2 (restored via object.__setattr__) + + Private (``init=False``) fields must travel in aux rather than as leaves so + that JAX round-trips (jit, vmap, grad, etc.) can restore them. They are + written back with ``object.__setattr__`` in ``unflatten``, which bypasses the + DataClass frozen guard safely because the object is still being initialised. + """ # Guard: If already registered (e.g., by __init_subclass__), skip re-registering if getattr(cls_, "_jax_tree_registered", False): return cls_ - static_fields = [f.name for f in fields(cls_) if f.metadata.get("jax_static")] - dynamic_fields = [f.name for f in fields(cls_) if not f.metadata.get("jax_static")] + all_fields = fields(cls_) + # Traced leaves: init=True and not jax_static + dynamic_fields = [ + f.name for f in all_fields if not f.metadata.get("jax_static") and f.init + ] + # Aux bucket 1: explicitly static (jax_static=True) and init=True + static_init_fields = [ + f.name for f in all_fields if f.metadata.get("jax_static") and f.init + ] + # Aux bucket 2: private (init=False), regardless of jax_static — must ride in aux + private_fields_names = [f.name for f in all_fields if not f.init] def flatten_with_keys(obj): keyed_leaves = [ (jax.tree_util.GetAttrKey(f), getattr(obj, f)) for f in dynamic_fields ] - aux = tuple(getattr(obj, f) for f in static_fields) + aux = ( + tuple(getattr(obj, f) for f in static_init_fields), + tuple(getattr(obj, f) for f in private_fields_names), + ) return keyed_leaves, aux def unflatten(aux, leaves): - kwargs = {**dict(zip(static_fields, aux)), **dict(zip(dynamic_fields, leaves))} - return cls_(**kwargs) + static_init_values, private_values = aux + init_kwargs = { + **dict(zip(static_init_fields, static_init_values)), + **dict(zip(dynamic_fields, leaves)), + } + obj = cls_(**init_kwargs) + # Stamp private fields back after construction. object.__setattr__ bypasses + # DataClass.__setattr__'s frozen guard, which is safe here because we are + # restoring the exact values the object was flattened from. + for name, value in zip(private_fields_names, private_values): + object.__setattr__(obj, name, value) + return obj jax.tree_util.register_pytree_with_keys(cls_, flatten_with_keys, unflatten) cls_._jax_tree_registered = True # ty:ignore[unresolved-attribute] diff --git a/src/drinx/visualize.py b/src/drinx/visualize.py index 8b0674c..0735499 100644 --- a/src/drinx/visualize.py +++ b/src/drinx/visualize.py @@ -60,9 +60,7 @@ def visualize_leaf(val: int | float | complex | bool | np.ndarray | jax.Array) - # 2. Build compact dtype string (NumPy's dtype.kind already returns 'f', 'i', 'u', 'c', 'b') dtype_str = _dtype_str(dtype) - prefix = ( - f"{dtype_str}[{','.join(map(str, shape))}]" # ty:ignore[invalid-argument-type] - ) + prefix = f"{dtype_str}[{','.join(str(d) for d in shape)}]" # ty:ignore[no-matching-overload] # 3. Handle Tracers if is_traced(val):