Silence some pytype errors.

PiperOrigin-RevId: 510514175

Author: rchen152

dataclass_array/array_dataclass.py

@@ -422,7 +422,7 @@ def replace(self: _DcT, **kwargs: Any) -> _DcT:
     }
 
     # Create the new object
-    new_self = dataclasses.replace(self, **init_kwargs)
+    new_self = dataclasses.replace(self, **init_kwargs)  # pytype: disable=wrong-arg-types  # re-none
 
     # TODO(epot): Could try to unify logic bellow with `tree_unflatten`
 
@@ -677,7 +677,7 @@ def tree_flatten(self) -> tuple[list[DcOrArray], _TreeMetadata]:
         array_field_names=list(self._all_array_fields.keys()),
         non_array_field_kwargs={
             f.name: getattr(self, f.name)
-            for f in dataclasses.fields(self)
+            for f in dataclasses.fields(self)  # pytype: disable=wrong-arg-types  # re-none
             if f.name not in self._all_array_fields  # pylint: disable=unsupported-membership-test
         },
     )
@@ -698,7 +698,7 @@ def tree_unflatten(
     )
     init_fields = {}
     non_init_fields = {}
-    fields = {f.name: f for f in dataclasses.fields(cls)}
+    fields = {f.name: f for f in dataclasses.fields(cls)}  # pytype: disable=wrong-arg-types  # re-none
     for k, v in metadata.non_array_field_kwargs.items():
       if fields[k].init:
         init_fields[k] = v

dataclass_array/utils/tree_utils_test.py

@@ -66,7 +66,7 @@ def add_fn(x, y):
 
   x0 = {'a': [1, 2]}
   x1 = {'a': [10, 20]}
-  assert tree_utils.tree_map(add_fn, x0, x1) == {'a': [11, 22]}
+  assert tree_utils.tree_map(add_fn, x0, x1) == {'a': [11, 22]}  # pytype: disable=wrong-arg-types  # re-none
 
 
 def test_tree_map_chex():
@@ -79,4 +79,4 @@ def add_fn(x, y):
 
   x0 = {'a': A(x=1, y=2)}
   x1 = {'a': A(x=10, y=20)}
-  assert tree_utils.tree_map(add_fn, x0, x1) == {'a': A(x=11, y=22)}
+  assert tree_utils.tree_map(add_fn, x0, x1) == {'a': A(x=11, y=22)}  # pytype: disable=wrong-arg-types  # re-none