Skip to content

updating dataclasses.replace and register_jax_tree #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
renaissancenerd wants to merge 1 commit into from
Closed

updating dataclasses.replace and register_jax_tree #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 59 additions & 4 deletions src/drinx/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment on lines +45 to +68

@coderabbitai coderabbitai Bot Apr 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

_dataclass_replace currently both exposes and drops private fields.

known_names lets init=False names through, so updated_copy(_cache=...) / aset("_cache", ...) now bypass the existing guard that tests/test_base.py:1302-1311 and tests/test_base.py:1400-1410 assert. At the same time, only f.name in changes gets stamped back, so updated_copy(x=...) still resets any existing private state instead of preserving it.

Suggested fix 
     all_fields = dataclasses.fields(obj)
-    known_names = {f.name for f in all_fields}
-    unknown = set(changes) - known_names
+    fields_by_name = {f.name: f for f in all_fields}
+    unknown = set(changes) - fields_by_name.keys()
     if unknown:
         raise TypeError(f"_dataclass_replace() got unexpected field names: {unknown!r}")
+
+    non_init_updates = {name for name in changes if not fields_by_name[name].init}
+    if non_init_updates:
+        raise TypeError(
+            f"_dataclass_replace() cannot update init=False fields: {non_init_updates!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
-    }
+    non_init_values = {f.name: current_values[f.name] for f in all_fields if not f.init}
 
     new_obj = type(obj)(**init_kwargs)
 
-    # Stamp non-init fields that were explicitly changed
-    for name, value in non_init_overrides.items():
+    # Preserve current non-init fields after construction
+    for name, value in non_init_values.items():
         object.__setattr__(new_obj, name, value)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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)
all_fields = dataclasses.fields(obj)
fields_by_name = {f.name: f for f in all_fields}
unknown = set(changes) - fields_by_name.keys()
if unknown:
raise TypeError(f"_dataclass_replace() got unexpected field names: {unknown!r}")
non_init_updates = {name for name in changes if not fields_by_name[name].init}
if non_init_updates:
raise TypeError(
f"_dataclass_replace() cannot update init=False fields: {non_init_updates!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_values = {f.name: current_values[f.name] for f in all_fields if not f.init}
new_obj = type(obj)(**init_kwargs)
# Preserve current non-init fields after construction
for name, value in non_init_values.items():
object.__setattr__(new_obj, name, value)
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/drinx/base.py` around lines 45 - 68, _dataclass_replace currently allows
changing private (underscore-prefixed) fields and also drops existing non-init
fields that aren't in changes; fix it by forbidding updates to private fields
and by restoring all non-init fields from the original object. Specifically: in
_dataclass_replace, compute private_names = {f.name for f in all_fields if
f.name.startswith("_")} and if private_names & set(changes): raise TypeError
(reject updates to private fields), keep known_names logic but validate against
public names only; build current_values as you do, then after creating new_obj,
stamp back every non-init field from current_values (use all non-init f.name),
and finally overwrite those non-init public fields with values from changes if
present (i.e., ensure non-init values are preserved unless an allowed public
non-init field was explicitly changed). Use the symbols _dataclass_replace,
known_names, current_values, non_init_overrides (or replace it with the new
logic) to locate and update the code.


return new_obj


@dataclass_transform(
field_specifiers=(
orig_field,
Expand Down Expand Up @@ -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"):
Expand Down Expand Up @@ -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")
Expand Down
45 changes: 39 additions & 6 deletions src/drinx/transform.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +29 to +63

@coderabbitai coderabbitai Bot Apr 16, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

🧩 Analysis chain

🌐 Web query:

In JAX custom pytree registration (register_pytree_with_keys), is aux_datatreated as static tree metadata, and should array-valued data that must participate injit/vmap/tree_mapbe returned as children/leaves rather than inaux_data?

💡 Result:

Yes, in JAX custom pytree registration with register_pytree_with_keys (and similarly register_pytree_node), aux_data is treated as static tree metadata. It is stored in the treedef, must be hashable and immutable, and is used for structure reconstruction during unflatten. It does not participate in transformations like jit, vmap, or tree_map—those operate on the dynamic children/leaves. Array-valued data that must participate in jit (i.e., be traced and potentially compiled differently based on value), vmap (vectorized), or tree_map (mapped over) should be returned as children/leaves, not in aux_data. Including arrays in aux_data leads to unsupported, poorly-defined behavior, such as unexpected tracing under jit (e.g., jnp arrays may become Traced unexpectedly, while np arrays do not). Official JAX documentation examples explicitly state: "children must contain arrays & pytrees" and "aux_data must contain static, hashable data." StackOverflow answer citing JAX docs confirms: "aux_data in a pytree flattening must contain static elements... Neither np.ndarray nor jax.Array satisfy this, so they should not be included in aux_data."

Citations:

🏁 Script executed:

#!/bin/bash
# Find definitions of private_field and static_private_field
rg "def (private_field|static_private_field)" -A 5

Repository: ymahlau/drinx

Length of output: 668

🏁 Script executed:

#!/bin/bash
# Search for uses of jax_static metadata in the fields
rg "jax_static" -B 2 -A 2 | head -100

Repository: ymahlau/drinx

Length of output: 4940

🏁 Script executed:

#!/bin/bash
# Check the structure of fields and metadata in base.py
fd "base.py" | xargs cat -n | head -150

Repository: ymahlau/drinx

Length of output: 6448

Split private fields by jax_static in pytree flatten/unflatten logic.

private_field() and static_private_field() are distinct APIs: the former defaults to dynamic (non-static) and the latter is always static. The current implementation collapses all init=False fields into pytree aux regardless of jax_static, placing non-static array-valued private fields into JAX's static treedef metadata where they cannot participate in transformations like jit, vmap, or tree_map.

Keep init=False and not jax_static fields as pytree leaves (via keyed_leaves), and only put init=False and jax_static values into aux. This aligns with the API design and JAX's requirement that aux_data contain only static, hashable metadata—not arrays or mutable data.

🧰 Tools
🪛 Ruff (0.15.10)

[warning] 54-54: zip() without an explicit strict= parameter

Add explicit value for parameter strict=

(B905)

[warning] 55-55: zip() without an explicit strict= parameter

Add explicit value for parameter strict=

(B905)

[warning] 61-61: zip() without an explicit strict= parameter

Add explicit value for parameter strict=

(B905)

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/drinx/transform.py` around lines 29 - 63, The flatten/unflatten currently
lumps all init=False fields into aux; change this so private_fields are split by
jax_static: compute private_dynamic = [f.name for f in all_fields if not f.init
and not f.metadata.get("jax_static")] and private_static = [f.name for f in
all_fields if not f.init and f.metadata.get("jax_static")]; then extend
dynamic_fields (used in keyed_leaves) with private_dynamic so non-static private
fields become leaves in flatten_with_keys, and only include static_private and
static_init_fields in aux (e.g., aux = (tuple(getattr(obj, f) for f in
static_init_fields), tuple(getattr(obj, f) for f in private_static))); finally
update unflatten to accept aux = (static_init_values, private_static_values),
build init_kwargs by zipping static_init_fields and dynamic_fields (which now
includes private_dynamic), construct obj = cls_(**init_kwargs), and then
object.__setattr__ only for names in private_static to restore static private
fields.


jax.tree_util.register_pytree_with_keys(cls_, flatten_with_keys, unflatten)
cls_._jax_tree_registered = True # ty:ignore[unresolved-attribute]
Expand Down
4 changes: 1 addition & 3 deletions src/drinx/visualize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading