Better fake array support

PiperOrigin-RevId: 595326519

Author: Conchylicultor

CHANGELOG.md

@@ -23,7 +23,7 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
-*   Fix `etree.spec_like` support
+*   Fix `etree.spec_like`, `jax.ShapeDtypeStruct`,... support
 
 ## [1.5.1] - 2023-08-30
 

dataclass_array/array_dataclass.py

@@ -1103,15 +1103,15 @@ def is_value_missing(self) -> bool:
       # In `jax/_src/api_util.py` for `flatten_axes`, jax set all values to a
       # dummy sentinel `object()` value.
       return True
-    elif isinstance(self.value, enp.ArraySpec):
-      # `etree.spec_like` compatibility
-      return True
     elif (
         isinstance(self.value, DataclassArray) and not self.value._array_fields  # pylint: disable=protected-access
     ):
       # Nested dataclass case (if all attributes are `object`, so no active
       # array fields)
       return True
+    elif enp.array_spec.is_fake_array(self.value):
+      # `etree.spec_like`, Flax summary, ShapeDtypeStruct, ... compatibility
+      return True
     return False
 
   @property