Add dca.typing alias for array types (alias of enp.typing)

Rational: This allow to use `dataclass_array` in a self-contained fashion, without importing 2 different packages
PiperOrigin-RevId: 469692004

Author: Conchylicultor

CHANGELOG.md

@@ -23,6 +23,9 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 ## [Unreleased]
 
+## [1.1.0] - 2022-08-15
+
+*   Added: Array types can be imported directly from `dataclass_array.typing`
 *   Added: Syntax to specify the shape of the DataclassArray (e.g. `MyRay['h
     w']`).
 *   Fixed: Correctly forward non-init fields in `.replace`, `tree_map`,
@@ -32,6 +35,7 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`):
 
 *   Initial release
 
-[Unreleased]: https://github.com/google-research/dataclass_array/compare/v1.0.0...HEAD
+[Unreleased]: https://github.com/google-research/dataclass_array/compare/v1.1.0...HEAD
+[1.1.0]: https://github.com/google-research/dataclass_array/compare/v1.0.0...v1.1.0
 [1.0.0]: https://github.com/google-research/dataclass_array/compare/v0.1.0...v1.0.0
 [0.1.0]: https://github.com/google-research/dataclass_array/releases/tag/v0.1.0

README.md

@@ -20,12 +20,14 @@ To view an example of dataclass arrays used in practice, see
 To create a `dca.DataclassArray`, take a frozen dataclass and:
 
 *   Inherit from `dca.DataclassArray`
-*   Annotate the fields with `etils.array_types` to specify the inner shape and
-    dtype of the array (see below for static or nested dataclass fields).
+*   Annotate the fields with `dataclass_array.typing` to specify the inner shape
+    and dtype of the array (see below for static or nested dataclass fields).
+    The array types are an alias from
+    [`etils.array_types`](https://github.com/google/etils/blob/main/etils/array_types/README.md).
 
 ```python
 import dataclass_array as dca
-from etils.array_types import FloatArray
+from dataclass_array.typing import FloatArray
 
 
 @dataclasses.dataclass(frozen=True)
@@ -145,8 +147,8 @@ For example, with `__matmul__(self, x: T) -> T`:
 (a, *x) @ (b, *x) -> Error: Incompatible a != b
 ```
 
-To test on Colab, see the `visu3d`
-dataclass [Colab tutorial](https://colab.research.google.com/github/google-research/visu3d/blob/main/docs/dataclass.ipynb).
+To test on Colab, see the `visu3d` dataclass
+[Colab tutorial](https://colab.research.google.com/github/google-research/visu3d/blob/main/docs/dataclass.ipynb).
 
 ## Motivating examples
 

dataclass_array/__init__.py

@@ -43,6 +43,6 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased
 # When changing this, also update the CHANGELOG.md
-__version__ = '1.0.0'
+__version__ = '1.1.0'
 
 del sys, pytest

dataclass_array/array_dataclass.py

@@ -23,6 +23,7 @@
 from dataclass_array import field_utils
 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.utils import np_utils
 from dataclass_array.utils import py_utils
@@ -31,7 +32,6 @@
 from etils import edc
 from etils import enp
 from etils import epy
-from etils.array_types import Array
 import numpy as np
 import typing_extensions
 from typing_extensions import Annotated, Literal, TypeAlias  # pylint: disable=g-multiple-import

dataclass_array/array_dataclass_test.py

@@ -20,9 +20,9 @@
 from typing import Any, Callable
 
 import dataclass_array as dca
+from dataclass_array.typing import FloatArray, IntArray, f32, i32  # pylint: disable=g-multiple-import
 from dataclass_array.typing import Shape  # pylint: disable=g-multiple-import
 from etils import enp
-from etils.array_types import FloatArray, IntArray, f32, i32  # pylint: disable=g-multiple-import
 import numpy as np
 import pytest
 import tensorflow as tf

dataclass_array/shape_parsing_test.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 from dataclass_array import shape_parsing
-from etils.array_types import f32
+from dataclass_array.typing import f32
 import pytest
 
 

dataclass_array/testing.py

@@ -21,7 +21,7 @@
 from typing import Any, Optional
 
 from dataclass_array import array_dataclass
-from etils.array_types import FloatArray  # pylint: disable=g-multiple-import
+from dataclass_array.typing import FloatArray  # pylint: disable=g-multiple-import
 from etils.etree import jax as etree
 from etils.etree import Tree
 import numpy as np

dataclass_array/type_parsing_test.py

@@ -21,7 +21,7 @@
 
 import dataclass_array as dca
 from dataclass_array import type_parsing
-from etils.array_types import f32, FloatArray  # pylint: disable=g-multiple-import
+from dataclass_array.typing import f32, FloatArray  # pylint: disable=g-multiple-import
 import pytest
 
 

dataclass_array/typing.py

@@ -19,11 +19,40 @@
 import typing
 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
 
+# ======== Array types (alias of `enp.typing`) ========
+
+ArrayAliasMeta = enp.typing.ArrayAliasMeta
+ArrayLike = enp.typing.ArrayLike
+
+Array = enp.typing.Array
+FloatArray = enp.typing.FloatArray
+IntArray = enp.typing.IntArray
+BoolArray = enp.typing.BoolArray
+StrArray = enp.typing.StrArray
+
+ui8 = enp.typing.ui8
+ui16 = enp.typing.ui16
+ui32 = enp.typing.ui32
+ui64 = enp.typing.ui64
+i8 = enp.typing.i8
+i16 = enp.typing.i16
+i32 = enp.typing.i32
+i64 = enp.typing.i64
+f16 = enp.typing.f16
+f32 = enp.typing.f32
+f64 = enp.typing.f64
+complex64 = enp.typing.complex64
+complex128 = enp.typing.complex128
+bool_ = enp.typing.bool_
+
+# ======== Dataclass array specific types ========
+
 TypeAlias = Any
 
 Shape = Tuple[int, ...]

dataclass_array/vectorization_test.py

@@ -20,10 +20,10 @@
 
 import dataclass_array as dca
 from dataclass_array import vectorization
+from dataclass_array.typing import FloatArray
 from dataclass_array.utils import inspect_utils
 from dataclass_array.utils import np_utils
 from etils import enp
-from etils.array_types import FloatArray
 import jax
 import pytest
 import tensorflow.experimental.numpy as tnp