Disable registration of dataclasses defined in __main__ with JAX tree…

… util.

This avoids pickling failures of the sort _pickle.PicklingError: Can't pickle <functools._lru_cache_wrapper object>: it's not the same object as register_dataclass_type_with_jax_tree_util.

PiperOrigin-RevId: 636154723

chex/_src/dataclass.py

@@ -228,7 +228,12 @@ def _getstate(self):
     # by JAX.
     #
     # See internal dataclass_test for unit tests demonstrating the problems.
-    register_dataclass_type_with_jax_tree_util(dcls)
+    # The registration below may result in pickling failures of the sort
+    # _pickle.PicklingError: Can't pickle <functools._lru_cache_wrapper object>:
+    # it's not the same object as register_dataclass_type_with_jax_tree_util
+    # for modules defined in __main__ so we disable registration in this case.
+    if dcls.__module__ != "__main__":
+      register_dataclass_type_with_jax_tree_util(dcls)
 
     # Patch __setstate__ to register the dataclass on deserialization.
     def _setstate(self, state):