Repository: patrick-kidger/equinox (python)

Issue:
Modules annotated with jaxtyping somtimes does not unpickle correctly
```python
import pickle

import equinox as eqx
from jax import Array
from jax import numpy as jnp
from jax import tree_util as jtu
from jaxtyping import Float


class A(eqx.Module):
    x: Float[Array, ""]  # works if replaced with just Array


def main():
    b = A(jnp.array(1.0))

    tree = jtu.tree_structure(b)
    b = tree.unflatten([jnp.array(2.0)])  # works if remove this line

    b2 = pickle.loads(pickle.dumps(b))

    # PyTreeDef(CustomNode(A[('x',), (), ()], [*]))
    print(jtu.tree_structure(b2))
    # PyTreeDef(CustomNode(A[('x',), (), ()], [*]))
    print(tree)

    # AssertionError
    assert jtu.tree_structure(b2) == tree
```

```
jax 0.5.3
equinox 0.11.12
jaxtyping 0.2.38
```


Pull request description:
Fix round trip flattening + pickling
Fixes #992

Patch:
diff --git a/equinox/_module.py b/equinox/_module.py
index 688f35c..94ec55b 100644
--- a/equinox/_module.py
+++ b/equinox/_module.py
@@ -881,7 +881,7 @@ def _flatten_module(module: "Module", with_keys: bool):
                 dynamic_field_values.append(value)
     sentinel = object()
     for name in _wrapper_field_names:
-        value = getattr(module, name, sentinel)
+        value = module.__dict__.get(name, sentinel)
         if value is not sentinel:
             wrapper_field_names.append(name)
             wrapper_field_values.append(value)


Test patch:
diff --git a/tests/test_module.py b/tests/test_module.py
index ad37eee..b9a7471 100644
--- a/tests/test_module.py
+++ b/tests/test_module.py
@@ -2,6 +2,7 @@ import abc
 import dataclasses
 import functools as ft
 import inspect
+import pickle
 from collections.abc import Callable
 from dataclasses import InitVar
 from typing import Any
@@ -12,6 +13,7 @@ import jax
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import pytest
+from jaxtyping import Array, Float
 
 from .helpers import tree_allclose
 
@@ -1241,3 +1243,17 @@ def test_converter_subclass_post_init():
     ax, bx = xs
     assert type(ax) is int
     assert type(bx) is int
+
+
+class _RoundTrip(eqx.Module):
+    foo: Float[Array, ""]
+
+
+# https://github.com/patrick-kidger/equinox/issues/992
+def test_tree_roundtrip_serialise_roundtrip():
+    x = _RoundTrip(jnp.array(1.0))
+    leaves, tree = jtu.tree_flatten(x)
+    x = jtu.tree_unflatten(tree, leaves)
+    x = pickle.loads(pickle.dumps(x))
+    tree2 = jtu.tree_structure(x)
+    assert tree == tree2