Fix jax.tree_util.register_dataclass in older JAX versions.

PiperOrigin-RevId: 628149376
google · Apr 25, 2024 · ea3bcab · ea3bcab
1 parent b8ccb15
commit ea3bcab
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/flax/struct.py b/flax/struct.py
@@ -124,11 +124,14 @@ def replace(self, **updates):
   data_clz.replace = replace
 
   # Remove this guard once minimux JAX version is >0.4.26.
-  if hasattr(jax.tree_util, 'register_dataclass'):
-    jax.tree_util.register_dataclass(
-        data_clz, data_fields, meta_fields
-    )
-  else:
+  try:
+    if hasattr(jax.tree_util, 'register_dataclass'):
+      jax.tree_util.register_dataclass(
+          data_clz, data_fields, meta_fields
+      )
+    else:
+      raise NotImplementedError
+  except NotImplementedError:
 
     def iterate_clz(x):
       meta = tuple(getattr(x, name) for name in meta_fields)