Fix next release pytype errors

PiperOrigin-RevId: 544331428

Author: Conchylicultor

dataclass_array/array_dataclass.py

@@ -25,7 +25,7 @@
 from dataclass_array import shape_parsing
 from dataclass_array import type_parsing
 from dataclass_array.typing import Array
-from dataclass_array.typing import Axes, DcOrArray, DcOrArrayT, DTypeArg, DynamicShape, Shape  # pylint: disable=g-multiple-import
+from dataclass_array.typing import Axes, DTypeArg, DcOrArray, DcOrArrayT, DynamicShape, Shape  # pylint: disable=g-multiple-import,g-importing-member
 from dataclass_array.utils import np_utils
 from dataclass_array.utils import py_utils
 import einops
@@ -37,8 +37,6 @@
 import typing_extensions
 from typing_extensions import Annotated, Literal, TypeAlias  # pylint: disable=g-multiple-import
 
-if typing.TYPE_CHECKING:
-  import torch  # pytype: disable=import-error
 
 lazy = enp.lazy
 
@@ -152,15 +150,16 @@ def __getitem__(cls, spec):
     return Annotated[cls, field_utils.ShapeAnnotation(spec)]
 
 
-@typing_extensions.dataclass_transform(  # pytype: disable=not-supported-yet
-    kw_only_default=True,
-    # TODO(b/272524683):Restore field specifier
-    # field_specifiers=(
-    #     dataclasses.Field,
-    #     dataclasses.field,
-    #     array_field,
-    # ),
-)
+# TODO(epot): Restore once pytype support this
+# @typing_extensions.dataclass_transform(
+#     kw_only_default=True,
+#     # TODO(b/272524683):Restore field specifier
+#     # field_specifiers=(
+#     #     dataclasses.Field,
+#     #     dataclasses.field,
+#     #     array_field,
+#     # ),
+# )
 class DataclassArray(metaclass=MetaDataclassArray):
   """Dataclass which behaves like an array.
 

dataclass_array/typing.py

@@ -20,10 +20,9 @@
 from typing import Any, Optional, Tuple, Type, TypeVar, Union
 
 from etils import enp
-from etils.array_types import FloatArray
 
 if typing.TYPE_CHECKING:
-  from dataclass_array import array_dataclass
+  from dataclass_array import array_dataclass  # pylint: disable=g-bad-import-order,unused-import
 
 # ======== Array types (alias of `enp.typing`) ========
 

dataclass_array/utils/tree_utils.py

@@ -22,14 +22,19 @@
 from __future__ import annotations
 
 import collections.abc
-from typing import Callable, TypeVar
+from typing import Any, Callable, TypeVar
 
 from etils import epy
-from etils.etree.typing import Tree
-from typing_extensions import Unpack, TypeVarTuple  # pytype: disable=not-supported-yet  # pylint: disable=g-multiple-import
+from etils.etree.typing import Tree  # pylint: disable=g-importing-member
+from typing_extensions import TypeVarTuple, Unpack  # pytype: disable=not-supported-yet  # pylint: disable=g-multiple-import
 
-_InsT = TypeVarTuple('_InsT')
-_OutT = TypeVar('_OutT')
+# TODO(epot): Remove once pytype support `Unpack`
+del TypeVar, TypeVarTuple, Unpack
+Unpack = Any
+TypeVarTuple = Any
+
+_InsT = Any  # TypeVarTuple('_InsT')
+_OutT = Any  # TypeVar('_OutT')
 
 
 def tree_map(  # pylint: disable=redefined-builtin