Fix _ArrayField.inner_shape for scalar values in PyTorch contexts.

Corrected an edge case in `_ArrayField.inner_shape` where, when called on a scalar value (empty `full_shape`) within a PyTorch environment, it now returns an empty shape `()` to prevent shape validation errors, particularly when used inside `torch.vmap`.

PiperOrigin-RevId: 855886831

Author: The dataclass_array Authors

dataclass_array/array_dataclass.py

@@ -1122,11 +1122,9 @@ def inner_shape(self) -> Shape:
     """Returns the the static shape resolved for the current value."""
     # torch.func.vmap calls `tree_unflatten([0] * num_leaves)` internally,
     # messing up shape inference.
-    if (
-        enp.lazy.has_torch
-        and isinstance(self.value, int)
-        and self.value == 0
-    ):
+    if not self.full_shape and self.inner_shape_non_static:
+      if enp.lazy.has_torch:
+        return ()
       return ()
     if not self.inner_shape_non_static:
       return ()