Fix _ArrayField.inner_shape for scalar values in PyTorch contexts.

The dataclass_array Authors · The dataclass_array Authors · commit a21b89ba4554 · 2026-01-12T10:20:39.000-08:00
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: 855281432
diff --git a/.github/workflows/pytest_and_autopublish.yml b/.github/workflows/pytest_and_autopublish.yml
@@ -22,8 +22,8 @@ jobs:
     - run: pip --version
     # As changes can be pushed to both etils and visu3d, we install etils from `main` branch
     # If modifying this, also modify `pyproject.toml`
-    - run: pip install "etils[array_types,edc,enp,epath,epy,etree] @ git+https://github.com/google/etils"
-    - run: pip install -e .[dev]
+    - run: pip install --no-cache-dir "etils[array_types,edc,enp,epath,epy,etree] @ git+https://github.com/google/etils"
+    - run: pip install --no-cache-dir -e .[dev]
     - run: pip freeze
 
     # Run tests (in parallel)
diff --git a/dataclass_array/__init__.py b/dataclass_array/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/array_dataclass.py b/dataclass_array/array_dataclass.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -1120,6 +1120,14 @@ def full_shape(self) -> DcOrArrayT:
   @functools.cached_property
   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
+    ):
+      return ()
     if not self.inner_shape_non_static:
       return ()
     static_shape = self.full_shape[-len(self.inner_shape_non_static) :]
diff --git a/dataclass_array/array_dataclass_test.py b/dataclass_array/array_dataclass_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/conftest.py b/dataclass_array/conftest.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/field_utils.py b/dataclass_array/field_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/import_test.py b/dataclass_array/import_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/ops.py b/dataclass_array/ops.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/shape_parsing.py b/dataclass_array/shape_parsing.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/shape_parsing_test.py b/dataclass_array/shape_parsing_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/testing.py b/dataclass_array/testing.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/type_parsing.py b/dataclass_array/type_parsing.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/type_parsing_test.py b/dataclass_array/type_parsing_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/typing.py b/dataclass_array/typing.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/file_utils.py b/dataclass_array/utils/file_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/inspect_utils.py b/dataclass_array/utils/inspect_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/inspect_utils_test.py b/dataclass_array/utils/inspect_utils_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/np_utils.py b/dataclass_array/utils/np_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/np_utils_test.py b/dataclass_array/utils/np_utils_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/py_utils.py b/dataclass_array/utils/py_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/tree_utils.py b/dataclass_array/utils/tree_utils.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/utils/tree_utils_test.py b/dataclass_array/utils/tree_utils_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/vectorization.py b/dataclass_array/vectorization.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/dataclass_array/vectorization_test.py b/dataclass_array/vectorization_test.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
diff --git a/docs/conf.py b/docs/conf.py
@@ -1,4 +1,4 @@
-# Copyright 2024 The dataclass_array Authors.
+# Copyright 2025 The dataclass_array Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# Copyright 2024 The dataclass_array Authors.`
	`1`	`+# Copyright 2025 The dataclass_array Authors.`
`2`	`2`	`#`
`3`	`3`	`# Licensed under the Apache License, Version 2.0 (the "License");`
`4`	`4`	`# you may not use this file except in compliance with the License.`