Skip to content

Commit 70288f5

Browse files
Jammy2211claude
authored andcommitted
fix(light/linear): mark pytree_token as ephemeral via __getstate__/__setstate__
pytree_token is a process-local counter increment used as a stable hash/eq identity for LightProfileLinear instances across jax.jit flatten/unflatten. It is not meaningful state to persist — copying a number from one process's counter into another's does not preserve identity across processes. PyAutoFit's database serializer routes every numeric attribute through the Value SQL row (sa.Float column), so persisting pytree_token as an int and reading it back yields a float. Python >=3.12 strictly requires __hash__ to return int, so any dict keyed on a LightProfileLinear (e.g. the linear_light_profile_intensity_dict in autogalaxy/abstract_fit.py) raised TypeError on the visualization path after a fit was loaded via FitImagingAgg. Add __getstate__ that omits pytree_token and __setstate__ that assigns a fresh value if missing. PyAutoFit's Instance._from_object and Object.__call__ already check for these methods and honour them, so no PyAutoFit changes are required. The JAX-jit path uses register_model's attr_const which reads vars(self) directly and is unaffected. Fixes the smoke-test failure on autolens_workspace_test/main: scripts/database/scrape/general.py. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent f641d69 commit 70288f5

2 files changed

Lines changed: 66 additions & 0 deletions

File tree

autogalaxy/profiles/light/linear/abstract.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,19 @@ def __eq__(self, other):
6262
and self.pytree_token == other.pytree_token
6363
)
6464

65+
def __getstate__(self):
66+
# pytree_token is a process-local counter increment used only as a
67+
# stable hash/eq identity for jax.jit flatten/unflatten round-trips.
68+
# Persisting it would copy a number from one process's counter into
69+
# another's — the JAX path uses register_model's attr_const which
70+
# reads vars(self) directly and is unaffected by __getstate__.
71+
return {k: v for k, v in self.__dict__.items() if k != "pytree_token"}
72+
73+
def __setstate__(self, state):
74+
self.__dict__.update(state)
75+
if "pytree_token" not in state:
76+
self.pytree_token = next(LightProfileLinear._pytree_token_counter)
77+
6578
@property
6679
def regularization(self):
6780
return None

test_autogalaxy/profiles/light/linear/test_abstract.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,3 +98,56 @@ def test__lp_instance_from__returns_instance_with_correct_intensity():
9898
)
9999

100100
assert lp_non_linear.intensity == 3.0
101+
102+
103+
def test__pytree_token_is_int_and_unique():
104+
lp_0 = ag.lp_linear.Sersic()
105+
lp_1 = ag.lp_linear.Sersic()
106+
107+
assert isinstance(lp_0.pytree_token, int)
108+
assert isinstance(lp_1.pytree_token, int)
109+
assert lp_0.pytree_token != lp_1.pytree_token
110+
111+
assert isinstance(hash(lp_0), int)
112+
assert hash(lp_0) == hash(lp_0)
113+
assert hash(lp_0) != hash(lp_1)
114+
115+
116+
def test__getstate__omits_pytree_token():
117+
lp = ag.lp_linear.Sersic()
118+
state = lp.__getstate__()
119+
120+
assert "pytree_token" not in state
121+
assert "effective_radius" in state
122+
123+
124+
def test__setstate__assigns_fresh_pytree_token_when_missing():
125+
lp = ag.lp_linear.Sersic()
126+
state = lp.__getstate__()
127+
128+
restored = ag.lp_linear.Sersic.__new__(ag.lp_linear.Sersic)
129+
restored.__setstate__(state)
130+
131+
assert isinstance(restored.pytree_token, int)
132+
assert isinstance(hash(restored), int)
133+
134+
135+
def test__pickle_roundtrip_preserves_int_hash():
136+
import pickle
137+
138+
lp = ag.lp_linear.Sersic()
139+
restored = pickle.loads(pickle.dumps(lp))
140+
141+
assert isinstance(hash(restored), int)
142+
assert isinstance(restored.pytree_token, int)
143+
assert restored.effective_radius == lp.effective_radius
144+
145+
146+
def test__setstate__preserves_pytree_token_when_present():
147+
lp = ag.lp_linear.Sersic()
148+
state_with_token = dict(lp.__dict__)
149+
150+
restored = ag.lp_linear.Sersic.__new__(ag.lp_linear.Sersic)
151+
restored.__setstate__(state_with_token)
152+
153+
assert restored.pytree_token == lp.pytree_token

0 commit comments

Comments
 (0)