Code update

PiperOrigin-RevId: 941084039

Author: h-joo

dataclass_array/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.

dataclass_array/array_dataclass.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.
@@ -137,7 +137,7 @@ def array_field(
   # TODO(epot): Validate shape, dtype
   dca_field = _ArrayFieldMetadata(
       inner_shape_non_static=shape,
-      dtype=dtype,
+      dtype=dtype,  # pyrefly: ignore[bad-argument-type]
   )
   return dataclasses.field(**field_kwargs, metadata={_METADATA_KEY: dca_field})
 
@@ -148,7 +148,7 @@ class MetaDataclassArray(type):
   # TODO(b/204422756): We cannot use `__class_getitem__` due to b/204422756
   def __getitem__(cls, spec):
     # Not clear how this would interact if cls is also a `Generic`
-    return Annotated[cls, field_utils.ShapeAnnotation(spec)]
+    return Annotated[cls, field_utils.ShapeAnnotation(spec)]  # pyrefly: ignore[not-a-type]
 
 
 # TODO(epot): Restore once pytype support this
@@ -386,7 +386,7 @@ def __getitem__(self: _DcT, indices: _IndicesArg) -> _DcT:
     indices = _to_absolute_indices(indices, shape=self.shape)
     return self._map_field(
         array_fn=lambda f: f.value[indices],
-        dc_fn=lambda f: f.value[indices],
+        dc_fn=lambda f: f.value[indices],  # pyrefly: ignore[bad-argument-type]
     )
 
   # _DcT[n *d] -> Iterator[_DcT[*d]]
@@ -454,7 +454,7 @@ def fn(ray: Optional[dca.Ray] = None):
 
   def map_field(
       self: _DcT,
-      fn: Callable[[Array['*din']], Array['*dout']],
+      fn: Callable[[Array['*din']], Array['*dout']],  # pyrefly: ignore[not-a-type]
   ) -> _DcT:
     """Apply a transformation on all arrays from the fields."""
     return self._map_field(  # pylint: disable=protected-access
@@ -659,7 +659,7 @@ def _cast_field(f: _ArrayField) -> None:
           new_value = np_utils.asarray(
               f.value,
               xnp=xnp,
-              dtype=f.dtype,
+              dtype=f.dtype,  # pyrefly: ignore[bad-argument-type]
               cast_dtype=self.__dca_params__.cast_dtype,
           )
         self._setattr(f.name, new_value)
@@ -740,7 +740,7 @@ def _to_absolute_axis(self, axis: Axes) -> Axes:
   def _map_field(
       self: _DcT,
       *,
-      array_fn: Callable[[_ArrayField[Array['*din']]], Array['*dout']],
+      array_fn: Callable[[_ArrayField[Array['*din']]], Array['*dout']],  # pyrefly: ignore[not-a-type]
       dc_fn: Optional[Callable[[_ArrayField[_DcT]], _DcT]],
       _inplace: bool = False,
   ) -> _DcT:
@@ -758,7 +758,7 @@ def _map_field(
 
     def _apply_field_dn(f: _ArrayField):
       if f.is_dataclass:  # Recurse on dataclasses
-        return dc_fn(f)  # pylint: disable=protected-access
+        return dc_fn(f)  # pylint: disable=protected-access  # pyrefly: ignore[not-callable]
       else:
         return array_fn(f)
 
@@ -881,7 +881,7 @@ def _setattr(self, name: str, value: Any) -> None:
     """Like setattr, but support `frozen` dataclasses."""
     object.__setattr__(self, name, value)
 
-  def assert_same_xnp(self, x: Union[Array[...], DataclassArray]) -> None:
+  def assert_same_xnp(self, x: Union[Array[...], DataclassArray]) -> None:  # pyrefly: ignore[not-a-type]
     """Assert the given array is of the same type as the current object."""
     xnp = np_utils.get_xnp(x)
     if xnp is not self.xnp:
@@ -942,7 +942,7 @@ def _init_cls(self: DataclassArray) -> None:
     # is propagated by the various ops.
     dca_fields_metadata[_DUMMY_ARRAY_FIELD] = _ArrayFieldMetadata(  # pytype: disable=wrong-arg-types
         inner_shape_non_static=(),
-        dtype=np.float32,
+        dtype=np.float32,  # pyrefly: ignore[bad-argument-type]
     )
     default_dummy_array = np.zeros((), dtype=np.float32)
     _add_field_to_dataclass(
@@ -1032,7 +1032,7 @@ def _to_absolute_indices(indices: _Indices, *, shape: Shape) -> _Indices:
   start_elems = indices[:ellipsis_index]
   end_elems = indices[ellipsis_index + 1 :]
   ellipsis_replacement = [slice(None)] * (len(shape) - valid_count)
-  return (*start_elems, *ellipsis_replacement, *end_elems)
+  return (*start_elems, *ellipsis_replacement, *end_elems)  # pyrefly: ignore[bad-return]
 
 
 @dataclasses.dataclass(frozen=True)
@@ -1116,15 +1116,15 @@ def full_shape(self) -> DcOrArrayT:
     """Access the `host.<field-name>.shape`."""
     # TODO(b/198633198): We need to convert to tuple because TF evaluate
     # empty shapes to True `bool(shape) == True` when `shape=()`
-    return tuple(self.value.shape)
+    return tuple(self.value.shape)  # pyrefly: ignore[bad-return, missing-attribute]
 
   @functools.cached_property
   def inner_shape(self) -> Shape:
     """Returns the the static shape resolved for the current value."""
 
     if not self.inner_shape_non_static:
       return ()
-    static_shape = self.full_shape[-len(self.inner_shape_non_static) :]
+    static_shape = self.full_shape[-len(self.inner_shape_non_static) :]  # pyrefly: ignore[unsupported-operation]
 
     def err_msg() -> ValueError:
       return ValueError(
@@ -1177,8 +1177,8 @@ def host_shape(self) -> Shape:
     if not self.inner_shape_non_static:
       shape = self.full_shape
     else:
-      shape = self.full_shape[: -len(self.inner_shape_non_static)]
-    return shape
+      shape = self.full_shape[: -len(self.inner_shape_non_static)]  # pyrefly: ignore[unsupported-operation]
+    return shape  # pyrefly: ignore[bad-return]
 
   def assert_shape(self) -> None:
     if self.host_shape + self.inner_shape != self.full_shape:
@@ -1191,7 +1191,7 @@ def broadcast_to(self, shape: Shape) -> DcOrArrayT:
     """Broadcast the host_shape."""
     final_shape = shape + self.inner_shape
     if self.is_dataclass:
-      return self.value.broadcast_to(final_shape)
+      return self.value.broadcast_to(final_shape)  # pyrefly: ignore[missing-attribute]
     else:
       return self.xnp.broadcast_to(self.value, final_shape)
 
@@ -1214,7 +1214,7 @@ def _is_called_inside_torch_vmap() -> bool:
 
 def _make_field_metadata(
     field: dataclasses.Field[Any],
-    hints: dict[str, TypeAlias],
+    hints: dict[str, TypeAlias],  # pyrefly: ignore[invalid-annotation]
 ) -> Optional[_ArrayFieldMetadata]:
   """Make the array field class."""
   # TODO(epot): One possible confusion is if user define
@@ -1233,7 +1233,7 @@ def _make_field_metadata(
   return field_metadata
 
 
-def _type_to_field_metadata(hint: TypeAlias) -> Optional[_ArrayFieldMetadata]:
+def _type_to_field_metadata(hint: TypeAlias) -> Optional[_ArrayFieldMetadata]:  # pyrefly: ignore[invalid-annotation]
   """Converts type hint to extract `inner_shape`, `dtype`."""
   array_type = type_parsing.get_array_type(hint)
 

dataclass_array/array_dataclass_test.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.
@@ -50,8 +50,8 @@ def assert_val(
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
 class Point(dca.DataclassArray):
-  x: f32['*shape']
-  y: f32['*shape']
+  x: f32['*shape']  # pyrefly: ignore[not-a-type]
+  y: f32['*shape']  # pyrefly: ignore[not-a-type]
 
   @staticmethod
   def make(shape: Shape, xnp: enp.NpModule) -> Point:
@@ -77,8 +77,8 @@ def assert_val(p: Point, shape: Shape, xnp: enp.NpModule = None):
 
 @dca.dataclass_array(broadcast=True, cast_dtype=True)
 class Isometrie(dca.DataclassArray):
-  r: f32['... 3 3']
-  t: i32[..., 2]
+  r: f32['... 3 3']  # pyrefly: ignore[not-a-type]
+  t: i32[..., 2]  # pyrefly: ignore[not-a-type]
 
   @staticmethod
   def make(shape: Shape, xnp: enp.NpModule) -> Isometrie:
@@ -198,8 +198,8 @@ class WithStatic(dca.DataclassArray):
 
   static: str
   nested_static: NestedOnlyStatic
-  x: f32['... 3']
-  y: Any = dca.field(shape=(2, 2), dtype=np.float32)
+  x: f32['... 3']  # pyrefly: ignore[not-a-type]
+  y: Any = dca.field(shape=(2, 2), dtype=np.float32)  # pyrefly: ignore[bad-argument-type]
 
   @staticmethod
   def make(shape: Shape, xnp: enp.NpModule) -> WithStatic:
@@ -410,7 +410,7 @@ def test_wrong_input_type():
   @dca.dataclass_array(broadcast=True, cast_dtype=True)
   class PointWrapper(dca.DataclassArray):
     pts: Point
-    rgb: f32['*shape 3']
+    rgb: f32['*shape 3']  # pyrefly: ignore[not-a-type]
 
   # ndarray instead of DataclassArray
   with pytest.raises(TypeError, match='Invalid PointWrapper.pts:'):
@@ -646,8 +646,8 @@ def fn(p: WithStatic) -> WithStatic:
 @enp.testing.parametrize_xnp()
 def test_dataclass_params_no_cast(xnp: enp.NpModule):
   class PointNoCast(dca.DataclassArray):
-    x: FloatArray['*shape']
-    y: IntArray['*shape']
+    x: FloatArray['*shape']  # pyrefly: ignore[not-a-type]
+    y: IntArray['*shape']  # pyrefly: ignore[not-a-type]
 
   with pytest.raises(ValueError, match='Cannot cast float16'):
     PointNoCast(
@@ -668,8 +668,8 @@ class PointNoCast(dca.DataclassArray):
 def test_dataclass_params_no_list(xnp: enp.NpModule):
   @dca.dataclass_array(cast_list=False)
   class PointNoList(dca.DataclassArray):
-    x: FloatArray['*shape']
-    y: IntArray['*shape']
+    x: FloatArray['*shape']  # pyrefly: ignore[not-a-type]
+    y: IntArray['*shape']  # pyrefly: ignore[not-a-type]
 
   with pytest.raises(TypeError, match='Could not infer numpy module'):
     PointNoList(
@@ -681,8 +681,8 @@ class PointNoList(dca.DataclassArray):
 @enp.testing.parametrize_xnp()
 def test_dataclass_params_no_broadcast(xnp: enp.NpModule):
   class PointNoBroadcast(dca.DataclassArray):
-    x: FloatArray['*shape']
-    y: IntArray['*shape']
+    x: FloatArray['*shape']  # pyrefly: ignore[not-a-type]
+    y: IntArray['*shape']  # pyrefly: ignore[not-a-type]
 
   with pytest.raises(ValueError, match='Cannot broadcast'):
     PointNoBroadcast(
@@ -695,8 +695,8 @@ class PointNoBroadcast(dca.DataclassArray):
 @pytest.mark.parametrize('batch_shape', [(), (1, 3)])
 def test_dataclass_none_shape(xnp: enp.NpModule, batch_shape: Shape):
   class PointDynamicShape(dca.DataclassArray):
-    x: FloatArray[..., None, None]
-    y: IntArray['... 3 _']
+    x: FloatArray[..., None, None]  # pyrefly: ignore[not-a-type]
+    y: IntArray['... 3 _']  # pyrefly: ignore[not-a-type]
 
   p = PointDynamicShape(
       x=xnp.zeros(batch_shape + (2, 3), dtype=xnp.float32),
@@ -751,8 +751,8 @@ class PointDynamicShape(dca.DataclassArray):
 @pytest.mark.parametrize('batch_shape', [(1,), (3,)])
 def test_concatenate(xnp: enp.NpModule, batch_shape: Shape):
   class TestConcatenateClass(dca.DataclassArray):
-    x: FloatArray['*shape 3']
-    y: FloatArray['*shape']
+    x: FloatArray['*shape 3']  # pyrefly: ignore[not-a-type]
+    y: FloatArray['*shape']  # pyrefly: ignore[not-a-type]
 
   p = TestConcatenateClass(
       x=xnp.zeros(batch_shape + (3,), dtype=xnp.float32),

dataclass_array/conftest.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.

dataclass_array/field_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.

dataclass_array/import_test.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.
@@ -27,7 +27,7 @@
 
 
 class A(dca.DataclassArray):
-  x: dca.typing.f32['*s']
+  x: dca.typing.f32['*s']  # pyrefly: ignore[not-a-type]
 
 
 def test_lazy():

dataclass_array/ops.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.
@@ -36,7 +36,7 @@ def _ops_base(
             int,
             Any,  # array_dataclass._ArrayField[Array['*din']],
         ],
-        Array['*dout'],
+        Array['*dout'],  # pyrefly: ignore[not-a-type]
     ],
     dc_fn: Optional[
         Callable[
@@ -73,7 +73,7 @@ def _ops_base(
   xnp = first_arr.xnp
   # If axis < 0, normalize the axis such as the last axis is before the inner
   # shape
-  axis = np_utils.to_absolute_axis(axis, ndim=first_arr.ndim + 1)
+  axis = np_utils.to_absolute_axis(axis, ndim=first_arr.ndim + 1)  # pyrefly: ignore[bad-assignment]
 
   # Iterating over only the fields of the `first_arr` will skip optional fields
   # if those are not set in `first_arr`, even if they are present in others.
@@ -82,7 +82,7 @@ def _ops_base(
   # Similarly, static values will be the ones from the first element.
   merged_arr = first_arr._map_field(  # pylint: disable=protected-access
       array_fn=functools.partial(array_fn, xnp, axis),
-      dc_fn=functools.partial(dc_fn, xnp, axis),
+      dc_fn=functools.partial(dc_fn, xnp, axis),  # pyrefly: ignore[bad-argument-type]
   )
   return merged_arr
 

dataclass_array/shape_parsing.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.

dataclass_array/shape_parsing_test.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.

dataclass_array/testing.py

@@ -1,4 +1,4 @@
-# Copyright 2025 The dataclass_array Authors.
+# Copyright 2026 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.
@@ -30,8 +30,8 @@
 class Ray(array_dataclass.DataclassArray):
   """Dummy dataclass array for testing."""
 
-  pos: FloatArray[..., 3]
-  dir: FloatArray[..., 3]
+  pos: FloatArray[..., 3]  # pyrefly: ignore[not-a-type]
+  dir: FloatArray[..., 3]  # pyrefly: ignore[not-a-type]
 
 
 # TODO(epot): Should use `chex.assert_xyz` once dataclasses support DM `tree`