From fc7a2c081ea48380261486799734375b33e83d65 Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Sun, 23 Oct 2022 21:11:05 -0700
Subject: [PATCH 01/19] Initial tensorflow stubs

---
 stubs/tensorflow/METADATA.toml |   5 ++
 stubs/tensorflow/__init__.pyi  | 149 +++++++++++++++++++++++++++++++++
 stubs/tensorflow/dtypes.pyi    |  52 ++++++++++++
 stubs/tensorflow/math.pyi      |  13 +++
 stubs/tensorflow/sparse.pyi    |  24 ++++++
 5 files changed, 243 insertions(+)
 create mode 100644 stubs/tensorflow/METADATA.toml
 create mode 100644 stubs/tensorflow/__init__.pyi
 create mode 100644 stubs/tensorflow/dtypes.pyi
 create mode 100644 stubs/tensorflow/math.pyi
 create mode 100644 stubs/tensorflow/sparse.pyi

diff --git a/stubs/tensorflow/METADATA.toml b/stubs/tensorflow/METADATA.toml
new file mode 100644
index 000000000000..4e65b63c3190
--- /dev/null
+++ b/stubs/tensorflow/METADATA.toml
@@ -0,0 +1,5 @@
+version = "2.10.*"
+requires = ["numpy"]
+
+[tool.stubtest]
+ignore_missing_stub = false
diff --git a/stubs/tensorflow/__init__.pyi b/stubs/tensorflow/__init__.pyi
new file mode 100644
index 000000000000..9f340975ba9c
--- /dev/null
+++ b/stubs/tensorflow/__init__.pyi
@@ -0,0 +1,149 @@
+from _typeshed import Incomplete
+from builtins import bool as _bool
+from contextlib import contextmanager
+from typing import Any, Iterable, Iterator, NoReturn, Sequence, TypeVar, overload
+from typing_extensions import ParamSpec, Self
+
+import numpy as np
+from tensorflow.dtypes import *
+
+# Most tf.math functions are exported as tf, but sadly not all are.
+from tensorflow.math import abs as abs
+from tensorflow.sparse import SparseTensor
+
+# Tensors ideally should be a generic type, but properly typing data type/shape
+# will be a lot of work. Until we have good non-generic tensorflow stubs,
+# we will skip making Tensor generic. Also good type hints for shapes will
+# run quickly into many places where type system is not strong enough today.
+# So shape typing is probably not worth doing anytime soon.
+_SliceT = int | slice | None
+
+_R = TypeVar("_R")
+_P = ParamSpec("_P")
+
+_FloatDataSequence = Sequence[float] | Sequence[_FloatDataSequence]
+_StrDataSequence = Sequence[str] | Sequence[_StrDataSequence]
+_ScalarTensorCompatible = Tensor | str | float | np.ndarray[Any, Any] | np.number[Any]
+_TensorCompatible = _ScalarTensorCompatible | Sequence[_TensorCompatible]
+_ShapeLike = TensorShape | Iterable[_ScalarTensorCompatible | None] | int | Tensor
+_DTypeLike = DType | str | np.dtype[Any]
+
+class Tensor:
+    def consumers(self) -> list[Incomplete]: ...
+    @property
+    def shape(self) -> TensorShape: ...
+    def get_shape(self) -> TensorShape: ...
+    @property
+    def dtype(self) -> DType: ...
+    @property
+    def graph(self) -> Graph: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def op(self) -> Operation: ...
+    def numpy(self) -> np.ndarray[Any, Any]: ...
+    def __int__(self) -> int: ...
+    def __abs__(self) -> Tensor: ...
+    def __add__(self, other: _TensorCompatible) -> Tensor: ...
+    def __radd__(self, other: _TensorCompatible) -> Tensor: ...
+    def __sub__(self, other: _TensorCompatible) -> Tensor: ...
+    def __rsub__(self, other: _TensorCompatible) -> Tensor: ...
+    def __mul__(self, other: _TensorCompatible) -> Tensor: ...
+    def __rmul__(self, other: _TensorCompatible) -> Tensor: ...
+    def __pow__(self, other: _TensorCompatible) -> Tensor: ...
+    def __matmul__(self, other: _TensorCompatible) -> Tensor: ...
+    def __rmatmul__(self, other: _TensorCompatible) -> Tensor: ...
+    def __floordiv__(self, other: _TensorCompatible) -> Tensor: ...
+    def __rfloordiv__(self, other: _TensorCompatible) -> Tensor: ...
+    def __truediv__(self, other: _TensorCompatible) -> Tensor: ...
+    def __rtruediv__(self, other: _TensorCompatible) -> Tensor: ...
+    def __neg__(self) -> Tensor: ...
+    def __and__(self, other: _TensorCompatible) -> Tensor: ...
+    def __rand__(self, other: _TensorCompatible) -> Tensor: ...
+    def __or__(self, other: _TensorCompatible) -> Tensor: ...
+    def __ror__(self, other: _TensorCompatible) -> Tensor: ...
+    def __eq__(self, other: _TensorCompatible) -> Tensor: ...
+    def __ne__(self, other: _TensorCompatible) -> Tensor: ...
+    def __ge__(self, other: _TensorCompatible) -> Tensor: ...
+    def __gt__(self, other: _TensorCompatible) -> Tensor: ...
+    def __le__(self, other: _TensorCompatible) -> Tensor: ...
+    def __lt__(self, other: _TensorCompatible) -> Tensor: ...
+    def __bool__(self) -> NoReturn: ...
+    def __getitem__(self, slice_spec: _SliceT | tuple[_SliceT, ...]) -> Tensor: ...
+    def __len__(self) -> int: ...
+    # This only works for rank 0 tensors.
+    def __index__(self) -> int: ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+class RaggedTensor:
+    def bounding_shape(
+        self, axis: _TensorCompatible | None = None, name: str | None = None, out_type: _DTypeLike | None = None
+    ) -> Tensor: ...
+    @classmethod
+    def from_sparse(
+        cls, st_input: SparseTensor, name: str | None = None, row_splits_dtype: _DTypeLike = int64
+    ) -> RaggedTensor: ...
+    def to_sparse(self, name: str | None = None) -> SparseTensor: ...
+    def to_tensor(
+        self, default_value: float | str | None = None, name: str | None = None, shape: _ShapeLike | None = None
+    ) -> Tensor: ...
+    def __add__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __radd__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __sub__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __mul__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __rmul__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __floordiv__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __truediv__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __getitem__(self, slice_spec: _SliceT | tuple[_SliceT, ...]) -> RaggedTensor: ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+class Operation:
+    @property
+    def inputs(self) -> list[Tensor]: ...
+    @property
+    def input_types(self) -> list[DType]: ...
+    @property
+    def outputs(self) -> list[Tensor]: ...
+    @property
+    def output_types(self) -> list[DType]: ...
+    @property
+    def device(self) -> str: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def type(self) -> str: ...
+    def __getitem__(self, slice_spec: int | slice | tuple[int | slice, ...]) -> Operation: ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+class TensorShape:
+    def __init__(self, dims: _ShapeLike): ...
+    @property
+    def rank(self) -> int: ...
+    def as_list(self) -> list[int | None]: ...
+    def assert_has_rank(self, rank: int) -> None: ...
+    def assert_is_compatible_with(self, other: Iterable[int | None]) -> None: ...
+    def __bool__(self) -> _bool: ...
+    @overload
+    def __getitem__(self, key: int) -> int | None: ...
+    @overload
+    def __getitem__(self, key: slice) -> TensorShape: ...
+    def __iter__(self) -> Iterator[int | None]: ...
+    def __len__(self) -> int: ...
+    def __add__(self, other: Iterable[int | None]) -> TensorShape: ...
+    def __radd__(self, other: Iterable[int | None]) -> TensorShape: ...
+    def __eq__(self, other: Iterable[int | None]) -> _bool: ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+class Graph:
+    def add_to_collection(self, name: str, value: object): ...
+    def add_to_collections(self, names: Iterable[str] | str, value: object): ...
+    @contextmanager
+    def as_default(self) -> Iterator[Self]: ...
+    def finalize(self) -> None: ...
+    def get_tensor_by_name(self, name: str) -> Tensor: ...
+    def get_operation_by_name(self, name: str) -> Operation: ...
+    def get_operations(self) -> list[Operation]: ...
+    def get_name_scope(self) -> str: ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/dtypes.pyi b/stubs/tensorflow/dtypes.pyi
new file mode 100644
index 000000000000..aec80b3be9ba
--- /dev/null
+++ b/stubs/tensorflow/dtypes.pyi
@@ -0,0 +1,52 @@
+from _typeshed import Incomplete
+from builtins import bool as _bool
+from typing import Any
+
+import numpy as np
+from tensorflow import _DTypeLike
+
+class DType:
+    @property
+    def name(self) -> str: ...
+    @property
+    def as_numpy_dtype(self) -> type[np.number[Any]]: ...
+    @property
+    def is_numpy_compatible(self) -> _bool: ...
+    @property
+    def is_bool(self) -> _bool: ...
+    @property
+    def is_floating(self) -> _bool: ...
+    @property
+    def is_integer(self) -> _bool: ...
+    @property
+    def is_quantized(self) -> _bool: ...
+    @property
+    def is_unsigned(self) -> _bool: ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+bool: DType = ...
+complex128: DType = ...
+complex64: DType = ...
+bfloat16: DType = ...
+float16: DType = ...
+half: DType = ...
+float32: DType = ...
+float64: DType = ...
+double: DType = ...
+int8: DType = ...
+int16: DType = ...
+int32: DType = ...
+int64: DType = ...
+uint8: DType = ...
+uint16: DType = ...
+uint32: DType = ...
+uint64: DType = ...
+qint8: DType = ...
+qint16: DType = ...
+qint32: DType = ...
+quint8: DType = ...
+quint16: DType = ...
+string: DType = ...
+
+def as_dtype(dtype: _DTypeLike) -> DType: ...
+def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/math.pyi b/stubs/tensorflow/math.pyi
new file mode 100644
index 000000000000..f97802a7a5da
--- /dev/null
+++ b/stubs/tensorflow/math.pyi
@@ -0,0 +1,13 @@
+from _typeshed import Incomplete
+from typing import overload
+
+from tensorflow import RaggedTensor, Tensor, _TensorCompatible
+from tensorflow.sparse import SparseTensor
+
+@overload
+def abs(x: _TensorCompatible, name: str | None = None) -> Tensor: ...
+@overload
+def abs(x: SparseTensor, name: str | None = None) -> SparseTensor: ...
+@overload
+def abs(x: RaggedTensor, name: str | None = None) -> RaggedTensor: ...
+def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/sparse.pyi b/stubs/tensorflow/sparse.pyi
new file mode 100644
index 000000000000..b5a17fa5d2ee
--- /dev/null
+++ b/stubs/tensorflow/sparse.pyi
@@ -0,0 +1,24 @@
+from _typeshed import Incomplete
+
+from tensorflow import Tensor, TensorShape, _TensorCompatible
+from tensorflow.dtypes import DType
+
+_SparseTensorCompatible = _TensorCompatible | SparseTensor
+
+class SparseTensor:
+    indices: Tensor
+    values: Tensor
+    dense_shape: Tensor
+    shape: TensorShape
+    dtype: DType
+    name: str
+    def __init__(self, indices: _TensorCompatible, values: _TensorCompatible, dense_shape: _TensorCompatible) -> None: ...
+    def get_shape(self) -> TensorShape: ...
+    # Many arithmetic operations are not directly supported. Some have alternatives like tf.sparse.add instead of +.
+    def __div__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
+    def __truediv__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
+    def __mul__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
+    def __rmul__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+def __getattr__(name: str) -> Incomplete: ...

From 29017d2c8bb98ec7e52f3b26009b57035b655b60 Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Sun, 23 Oct 2022 21:39:32 -0700
Subject: [PATCH 02/19] fix flake8 lints besides default values

---
 stubs/tensorflow/__init__.pyi | 22 ++++++++++------------
 stubs/tensorflow/sparse.pyi   |  3 ++-
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/stubs/tensorflow/__init__.pyi b/stubs/tensorflow/__init__.pyi
index 9f340975ba9c..6b983e4eda0e 100644
--- a/stubs/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/__init__.pyi
@@ -1,8 +1,9 @@
 from _typeshed import Incomplete
 from builtins import bool as _bool
+from collections import Iterable, Iterator, Sequence
 from contextlib import contextmanager
-from typing import Any, Iterable, Iterator, NoReturn, Sequence, TypeVar, overload
-from typing_extensions import ParamSpec, Self
+from typing import Any, NoReturn, overload
+from typing_extensions import Self, TypeAlias
 
 import numpy as np
 from tensorflow.dtypes import *
@@ -16,17 +17,14 @@ from tensorflow.sparse import SparseTensor
 # we will skip making Tensor generic. Also good type hints for shapes will
 # run quickly into many places where type system is not strong enough today.
 # So shape typing is probably not worth doing anytime soon.
-_SliceT = int | slice | None
+_SliceT: TypeAlias = int | slice | None
 
-_R = TypeVar("_R")
-_P = ParamSpec("_P")
-
-_FloatDataSequence = Sequence[float] | Sequence[_FloatDataSequence]
-_StrDataSequence = Sequence[str] | Sequence[_StrDataSequence]
-_ScalarTensorCompatible = Tensor | str | float | np.ndarray[Any, Any] | np.number[Any]
-_TensorCompatible = _ScalarTensorCompatible | Sequence[_TensorCompatible]
-_ShapeLike = TensorShape | Iterable[_ScalarTensorCompatible | None] | int | Tensor
-_DTypeLike = DType | str | np.dtype[Any]
+_FloatDataSequence: TypeAlias = Sequence[float] | Sequence[_FloatDataSequence]
+_StrDataSequence: TypeAlias = Sequence[str] | Sequence[_StrDataSequence]
+_ScalarTensorCompatible: TypeAlias = Tensor | str | float | np.ndarray[Any, Any] | np.number[Any]
+_TensorCompatible: TypeAlias = _ScalarTensorCompatible | Sequence[_TensorCompatible]
+_ShapeLike: TypeAlias = TensorShape | Iterable[_ScalarTensorCompatible | None] | int | Tensor
+_DTypeLike: TypeAlias = DType | str | np.dtype[Any]
 
 class Tensor:
     def consumers(self) -> list[Incomplete]: ...
diff --git a/stubs/tensorflow/sparse.pyi b/stubs/tensorflow/sparse.pyi
index b5a17fa5d2ee..e1562644d003 100644
--- a/stubs/tensorflow/sparse.pyi
+++ b/stubs/tensorflow/sparse.pyi
@@ -1,9 +1,10 @@
 from _typeshed import Incomplete
+from typing_extensions import TypeAlias
 
 from tensorflow import Tensor, TensorShape, _TensorCompatible
 from tensorflow.dtypes import DType
 
-_SparseTensorCompatible = _TensorCompatible | SparseTensor
+_SparseTensorCompatible: TypeAlias = _TensorCompatible | SparseTensor
 
 class SparseTensor:
     indices: Tensor

From 93d1db93c59b821e4460c2fb9639217d9af3be1d Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Sun, 23 Oct 2022 21:45:39 -0700
Subject: [PATCH 03/19] fix mypy primer

---
 stubs/tensorflow/{ => tensorflow}/__init__.pyi | 0
 stubs/tensorflow/{ => tensorflow}/dtypes.pyi   | 0
 stubs/tensorflow/{ => tensorflow}/math.pyi     | 0
 stubs/tensorflow/{ => tensorflow}/sparse.pyi   | 0
 4 files changed, 0 insertions(+), 0 deletions(-)
 rename stubs/tensorflow/{ => tensorflow}/__init__.pyi (100%)
 rename stubs/tensorflow/{ => tensorflow}/dtypes.pyi (100%)
 rename stubs/tensorflow/{ => tensorflow}/math.pyi (100%)
 rename stubs/tensorflow/{ => tensorflow}/sparse.pyi (100%)

diff --git a/stubs/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
similarity index 100%
rename from stubs/tensorflow/__init__.pyi
rename to stubs/tensorflow/tensorflow/__init__.pyi
diff --git a/stubs/tensorflow/dtypes.pyi b/stubs/tensorflow/tensorflow/dtypes.pyi
similarity index 100%
rename from stubs/tensorflow/dtypes.pyi
rename to stubs/tensorflow/tensorflow/dtypes.pyi
diff --git a/stubs/tensorflow/math.pyi b/stubs/tensorflow/tensorflow/math.pyi
similarity index 100%
rename from stubs/tensorflow/math.pyi
rename to stubs/tensorflow/tensorflow/math.pyi
diff --git a/stubs/tensorflow/sparse.pyi b/stubs/tensorflow/tensorflow/sparse.pyi
similarity index 100%
rename from stubs/tensorflow/sparse.pyi
rename to stubs/tensorflow/tensorflow/sparse.pyi

From 46369ede77130e78633974b5678602d40e416519 Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Mon, 24 Oct 2022 10:34:00 -0700
Subject: [PATCH 04/19] Remove default values for consistency with typeshed.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 10 ++++------
 stubs/tensorflow/tensorflow/math.pyi     |  6 +++---
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 6b983e4eda0e..0b95e2f6cb5f 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -75,15 +75,13 @@ class Tensor:
 
 class RaggedTensor:
     def bounding_shape(
-        self, axis: _TensorCompatible | None = None, name: str | None = None, out_type: _DTypeLike | None = None
+        self, axis: _TensorCompatible | None = ..., name: str | None = ..., out_type: _DTypeLike | None = ...
     ) -> Tensor: ...
     @classmethod
-    def from_sparse(
-        cls, st_input: SparseTensor, name: str | None = None, row_splits_dtype: _DTypeLike = int64
-    ) -> RaggedTensor: ...
-    def to_sparse(self, name: str | None = None) -> SparseTensor: ...
+    def from_sparse(cls, st_input: SparseTensor, name: str | None = ..., row_splits_dtype: _DTypeLike = ...) -> RaggedTensor: ...
+    def to_sparse(self, name: str | None = ...) -> SparseTensor: ...
     def to_tensor(
-        self, default_value: float | str | None = None, name: str | None = None, shape: _ShapeLike | None = None
+        self, default_value: float | str | None = ..., name: str | None = ..., shape: _ShapeLike | None = ...
     ) -> Tensor: ...
     def __add__(self, other: RaggedTensor | float) -> RaggedTensor: ...
     def __radd__(self, other: RaggedTensor | float) -> RaggedTensor: ...
diff --git a/stubs/tensorflow/tensorflow/math.pyi b/stubs/tensorflow/tensorflow/math.pyi
index f97802a7a5da..d767ba8fad35 100644
--- a/stubs/tensorflow/tensorflow/math.pyi
+++ b/stubs/tensorflow/tensorflow/math.pyi
@@ -5,9 +5,9 @@ from tensorflow import RaggedTensor, Tensor, _TensorCompatible
 from tensorflow.sparse import SparseTensor
 
 @overload
-def abs(x: _TensorCompatible, name: str | None = None) -> Tensor: ...
+def abs(x: _TensorCompatible, name: str | None = ...) -> Tensor: ...
 @overload
-def abs(x: SparseTensor, name: str | None = None) -> SparseTensor: ...
+def abs(x: SparseTensor, name: str | None = ...) -> SparseTensor: ...
 @overload
-def abs(x: RaggedTensor, name: str | None = None) -> RaggedTensor: ...
+def abs(x: RaggedTensor, name: str | None = ...) -> RaggedTensor: ...
 def __getattr__(name: str) -> Incomplete: ...

From 304756db9be97bc86b2911b63c34696d226eee4d Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Tue, 10 Jan 2023 19:45:46 -0800
Subject: [PATCH 05/19] Add default values and fix pyright errors.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 16 +++++++++-------
 stubs/tensorflow/tensorflow/math.pyi     |  6 +++---
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 0b95e2f6cb5f..9d85ab1590b7 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -75,13 +75,15 @@ class Tensor:
 
 class RaggedTensor:
     def bounding_shape(
-        self, axis: _TensorCompatible | None = ..., name: str | None = ..., out_type: _DTypeLike | None = ...
+        self, axis: _TensorCompatible | None = None, name: str | None = None, out_type: _DTypeLike | None = None
     ) -> Tensor: ...
     @classmethod
-    def from_sparse(cls, st_input: SparseTensor, name: str | None = ..., row_splits_dtype: _DTypeLike = ...) -> RaggedTensor: ...
-    def to_sparse(self, name: str | None = ...) -> SparseTensor: ...
+    def from_sparse(
+        cls, st_input: SparseTensor, name: str | None = None, row_splits_dtype: _DTypeLike = int64
+    ) -> RaggedTensor: ...
+    def to_sparse(self, name: str | None = None) -> SparseTensor: ...
     def to_tensor(
-        self, default_value: float | str | None = ..., name: str | None = ..., shape: _ShapeLike | None = ...
+        self, default_value: float | str | None = None, name: str | None = None, shape: _ShapeLike | None = None
     ) -> Tensor: ...
     def __add__(self, other: RaggedTensor | float) -> RaggedTensor: ...
     def __radd__(self, other: RaggedTensor | float) -> RaggedTensor: ...
@@ -112,7 +114,7 @@ class Operation:
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class TensorShape:
-    def __init__(self, dims: _ShapeLike): ...
+    def __init__(self, dims: _ShapeLike) -> None: ...
     @property
     def rank(self) -> int: ...
     def as_list(self) -> list[int | None]: ...
@@ -131,8 +133,8 @@ class TensorShape:
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class Graph:
-    def add_to_collection(self, name: str, value: object): ...
-    def add_to_collections(self, names: Iterable[str] | str, value: object): ...
+    def add_to_collection(self, name: str, value: object) -> None: ...
+    def add_to_collections(self, names: Iterable[str] | str, value: object) -> None: ...
     @contextmanager
     def as_default(self) -> Iterator[Self]: ...
     def finalize(self) -> None: ...
diff --git a/stubs/tensorflow/tensorflow/math.pyi b/stubs/tensorflow/tensorflow/math.pyi
index d767ba8fad35..f97802a7a5da 100644
--- a/stubs/tensorflow/tensorflow/math.pyi
+++ b/stubs/tensorflow/tensorflow/math.pyi
@@ -5,9 +5,9 @@ from tensorflow import RaggedTensor, Tensor, _TensorCompatible
 from tensorflow.sparse import SparseTensor
 
 @overload
-def abs(x: _TensorCompatible, name: str | None = ...) -> Tensor: ...
+def abs(x: _TensorCompatible, name: str | None = None) -> Tensor: ...
 @overload
-def abs(x: SparseTensor, name: str | None = ...) -> SparseTensor: ...
+def abs(x: SparseTensor, name: str | None = None) -> SparseTensor: ...
 @overload
-def abs(x: RaggedTensor, name: str | None = ...) -> RaggedTensor: ...
+def abs(x: RaggedTensor, name: str | None = None) -> RaggedTensor: ...
 def __getattr__(name: str) -> Incomplete: ...

From 5fc1ce4d35ad78f6b32daaeefa7b7ec7b1315234 Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Tue, 10 Jan 2023 20:17:00 -0800
Subject: [PATCH 06/19] Fix couple mypy errors.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 9d85ab1590b7..602c92ea68a9 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -1,9 +1,9 @@
-from _typeshed import Incomplete
+from _typeshed import Incomplete, Self
 from builtins import bool as _bool
-from collections import Iterable, Iterator, Sequence
+from collections.abc import Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from typing import Any, NoReturn, overload
-from typing_extensions import Self, TypeAlias
+from typing_extensions import TypeAlias
 
 import numpy as np
 from tensorflow.dtypes import *
@@ -60,8 +60,8 @@ class Tensor:
     def __rand__(self, other: _TensorCompatible) -> Tensor: ...
     def __or__(self, other: _TensorCompatible) -> Tensor: ...
     def __ror__(self, other: _TensorCompatible) -> Tensor: ...
-    def __eq__(self, other: _TensorCompatible) -> Tensor: ...
-    def __ne__(self, other: _TensorCompatible) -> Tensor: ...
+    def __eq__(self, other: _TensorCompatible) -> Tensor: ...  # type: ignore[override]
+    def __ne__(self, other: _TensorCompatible) -> Tensor: ...  # type: ignore[override]
     def __ge__(self, other: _TensorCompatible) -> Tensor: ...
     def __gt__(self, other: _TensorCompatible) -> Tensor: ...
     def __le__(self, other: _TensorCompatible) -> Tensor: ...
@@ -129,14 +129,13 @@ class TensorShape:
     def __len__(self) -> int: ...
     def __add__(self, other: Iterable[int | None]) -> TensorShape: ...
     def __radd__(self, other: Iterable[int | None]) -> TensorShape: ...
-    def __eq__(self, other: Iterable[int | None]) -> _bool: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class Graph:
     def add_to_collection(self, name: str, value: object) -> None: ...
     def add_to_collections(self, names: Iterable[str] | str, value: object) -> None: ...
     @contextmanager
-    def as_default(self) -> Iterator[Self]: ...
+    def as_default(self: Self) -> Iterator[Self]: ...
     def finalize(self) -> None: ...
     def get_tensor_by_name(self, name: str) -> Tensor: ...
     def get_operation_by_name(self, name: str) -> Operation: ...

From bb714f538583ee14c43da0624fa03a254b7b51aa Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Tue, 10 Jan 2023 21:26:51 -0800
Subject: [PATCH 07/19] Add ignore missing stub

---
 .flake8                                  | 2 +-
 stubs/tensorflow/METADATA.toml           | 5 +----
 stubs/tensorflow/tensorflow/__init__.pyi | 6 +++---
 3 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/.flake8 b/.flake8
index 3d2d0193a83a..02fe1e480501 100644
--- a/.flake8
+++ b/.flake8
@@ -26,7 +26,7 @@
 [flake8]
 per-file-ignores =
   *.py: E203, E301, E302, E305, E501
-  *.pyi: B, E301, E302, E305, E501, E701, E741, F401, F403, F405, F822, Y037
+  *.pyi: B, E301, E302, E305, E501, E701, E741, F401, F403, F405, F822, Y011, Y037
   # Since typing.pyi defines "overload" this is not recognized by flake8 as typing.overload.
   # Unfortunately, flake8 does not allow to "noqa" just a specific error inside the file itself.
   # https://github.com/PyCQA/flake8/issues/1079
diff --git a/stubs/tensorflow/METADATA.toml b/stubs/tensorflow/METADATA.toml
index 4e65b63c3190..219877d7648b 100644
--- a/stubs/tensorflow/METADATA.toml
+++ b/stubs/tensorflow/METADATA.toml
@@ -1,5 +1,2 @@
 version = "2.10.*"
-requires = ["numpy"]
-
-[tool.stubtest]
-ignore_missing_stub = false
+requires = ["numpy"]
\ No newline at end of file
diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 602c92ea68a9..b5b4e844650f 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -17,7 +17,7 @@ from tensorflow.sparse import SparseTensor
 # we will skip making Tensor generic. Also good type hints for shapes will
 # run quickly into many places where type system is not strong enough today.
 # So shape typing is probably not worth doing anytime soon.
-_SliceT: TypeAlias = int | slice | None
+_Slice: TypeAlias = int | slice | None
 
 _FloatDataSequence: TypeAlias = Sequence[float] | Sequence[_FloatDataSequence]
 _StrDataSequence: TypeAlias = Sequence[str] | Sequence[_StrDataSequence]
@@ -67,7 +67,7 @@ class Tensor:
     def __le__(self, other: _TensorCompatible) -> Tensor: ...
     def __lt__(self, other: _TensorCompatible) -> Tensor: ...
     def __bool__(self) -> NoReturn: ...
-    def __getitem__(self, slice_spec: _SliceT | tuple[_SliceT, ...]) -> Tensor: ...
+    def __getitem__(self, slice_spec: _Slice | tuple[_Slice, ...]) -> Tensor: ...
     def __len__(self) -> int: ...
     # This only works for rank 0 tensors.
     def __index__(self) -> int: ...
@@ -92,7 +92,7 @@ class RaggedTensor:
     def __rmul__(self, other: RaggedTensor | float) -> RaggedTensor: ...
     def __floordiv__(self, other: RaggedTensor | float) -> RaggedTensor: ...
     def __truediv__(self, other: RaggedTensor | float) -> RaggedTensor: ...
-    def __getitem__(self, slice_spec: _SliceT | tuple[_SliceT, ...]) -> RaggedTensor: ...
+    def __getitem__(self, slice_spec: _Slice | tuple[_Slice, ...]) -> RaggedTensor: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class Operation:

From df9576ef6472fa6d060acbd0016a038b3b89be03 Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Tue, 10 Jan 2023 23:01:18 -0800
Subject: [PATCH 08/19] Fix stubtest errors.

---
 .flake8                                       |  2 +-
 .../tensorflow/@tests/stubtest_allowlist.txt  | 22 +++++
 stubs/tensorflow/tensorflow/__init__.pyi      | 93 ++++++++++++++-----
 .../core/framework/variable_pb2.pyi           |  0
 stubs/tensorflow/tensorflow/dtypes.pyi        |  7 +-
 stubs/tensorflow/tensorflow/sparse.pyi        | 19 ++--
 6 files changed, 111 insertions(+), 32 deletions(-)
 create mode 100644 stubs/tensorflow/@tests/stubtest_allowlist.txt
 create mode 100644 stubs/tensorflow/tensorflow/core/framework/variable_pb2.pyi

diff --git a/.flake8 b/.flake8
index 02fe1e480501..3a94f54f3710 100644
--- a/.flake8
+++ b/.flake8
@@ -26,7 +26,7 @@
 [flake8]
 per-file-ignores =
   *.py: E203, E301, E302, E305, E501
-  *.pyi: B, E301, E302, E305, E501, E701, E741, F401, F403, F405, F822, Y011, Y037
+  *.pyi: B, E301, E302, E305, E501, E701, E741, F401, F403, F405, F822, Y011, Y015, Y037
   # Since typing.pyi defines "overload" this is not recognized by flake8 as typing.overload.
   # Unfortunately, flake8 does not allow to "noqa" just a specific error inside the file itself.
   # https://github.com/PyCQA/flake8/issues/1079
diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt
new file mode 100644
index 000000000000..e59da503fe04
--- /dev/null
+++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt
@@ -0,0 +1,22 @@
+# Some methods are dynamically patched onto to instances as they 
+# may depend on whether code is executed in graph/eager/v1/v2/etc.
+# Tensorflow supports multiple modes of execution which changes some
+# of the attributes/methods/even class hierachies.
+tensorflow.Tensor.__int__
+tensorflow.Tensor.numpy
+tensorflow.Tensor.__index__
+# Incomplete
+tensorflow.sparse.SparseTensor.__getattr__
+tensorflow.SparseTensor.__getattr__
+tensorflow.TensorShape.__getattr__
+tensorflow.dtypes.DType.__getattr__
+tensorflow.RaggedTensor.__getattr__
+tensorflow.DType.__getattr__
+tensorflow.Graph.__getattr__
+tensorflow.Operation.__getattr__
+tensorflow.Variable.__getattr__
+# Internal undocumented API
+tensorflow.RaggedTensor.__init__
+# Has an undocumented extra argument that tf.Variable which acts like subclass 
+# (by dynamically patching tf.Tensor methods) does not preserve.
+tensorflow.Tensor.__getitem__
\ No newline at end of file
diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index b5b4e844650f..2432959eae15 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -1,7 +1,9 @@
-from _typeshed import Incomplete, Self
+from _typeshed import Incomplete, Self, Unused
+from abc import ABCMeta
 from builtins import bool as _bool
-from collections.abc import Iterable, Iterator, Sequence
+from collections.abc import Callable, Iterable, Iterator, Sequence
 from contextlib import contextmanager
+from enum import Enum
 from typing import Any, NoReturn, overload
 from typing_extensions import TypeAlias
 
@@ -27,6 +29,7 @@ _ShapeLike: TypeAlias = TensorShape | Iterable[_ScalarTensorCompatible | None] |
 _DTypeLike: TypeAlias = DType | str | np.dtype[Any]
 
 class Tensor:
+    def __init__(self, op: Operation, value_index: int, dtype: DType) -> None: ...
     def consumers(self) -> list[Incomplete]: ...
     @property
     def shape(self) -> TensorShape: ...
@@ -41,7 +44,7 @@ class Tensor:
     def op(self) -> Operation: ...
     def numpy(self) -> np.ndarray[Any, Any]: ...
     def __int__(self) -> int: ...
-    def __abs__(self) -> Tensor: ...
+    def __abs__(self, name: str | None = None) -> Tensor: ...
     def __add__(self, other: _TensorCompatible) -> Tensor: ...
     def __radd__(self, other: _TensorCompatible) -> Tensor: ...
     def __sub__(self, other: _TensorCompatible) -> Tensor: ...
@@ -55,17 +58,17 @@ class Tensor:
     def __rfloordiv__(self, other: _TensorCompatible) -> Tensor: ...
     def __truediv__(self, other: _TensorCompatible) -> Tensor: ...
     def __rtruediv__(self, other: _TensorCompatible) -> Tensor: ...
-    def __neg__(self) -> Tensor: ...
+    def __neg__(self, name: str | None = None) -> Tensor: ...
     def __and__(self, other: _TensorCompatible) -> Tensor: ...
     def __rand__(self, other: _TensorCompatible) -> Tensor: ...
     def __or__(self, other: _TensorCompatible) -> Tensor: ...
     def __ror__(self, other: _TensorCompatible) -> Tensor: ...
     def __eq__(self, other: _TensorCompatible) -> Tensor: ...  # type: ignore[override]
     def __ne__(self, other: _TensorCompatible) -> Tensor: ...  # type: ignore[override]
-    def __ge__(self, other: _TensorCompatible) -> Tensor: ...
-    def __gt__(self, other: _TensorCompatible) -> Tensor: ...
-    def __le__(self, other: _TensorCompatible) -> Tensor: ...
-    def __lt__(self, other: _TensorCompatible) -> Tensor: ...
+    def __ge__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
+    def __gt__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
+    def __le__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
+    def __lt__(self, other: _TensorCompatible, name: str | None = None) -> Tensor: ...
     def __bool__(self) -> NoReturn: ...
     def __getitem__(self, slice_spec: _Slice | tuple[_Slice, ...]) -> Tensor: ...
     def __len__(self) -> int: ...
@@ -73,7 +76,46 @@ class Tensor:
     def __index__(self) -> int: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
-class RaggedTensor:
+class VariableSynchronization(Enum):
+    AUTO = 0
+    NONE = 1
+    ON_WRITE = 2
+    ON_READ = 3
+
+class VariableAggregation(Enum):
+    AUTO = 0
+    NONE = 1
+    ON_WRITE = 2
+    ON_READ = 3
+
+class _VariableMetaclass(type): ...
+
+# Variable class in intent/documentation is a Tensor. In implementation there's
+# TODO comment to make it Tensor. It is not actually Tensor type wise, but even
+# dynamically patches on most methods of tf.Tensor
+# https://github.com/tensorflow/tensorflow/blob/9524a636cae9ae3f0554203c1ba7ee29c85fcf12/tensorflow/python/ops/variables.py#L1086.
+class Variable(Tensor, metaclass=_VariableMetaclass):
+    def __init__(
+        self,
+        initial_value: Tensor | Callable[[], Tensor] | None = None,
+        trainable: _bool | None = None,
+        validate_shape: _bool = True,
+        # Valid non-None values are deprecated.
+        caching_device: None = None,
+        name: str | None = None,
+        # Real type is VariableDef protobuf type. Can be added after adding script
+        # to generate tensorflow protobuf stubs with mypy-protobuf.
+        variable_def: Incomplete | None = None,
+        dtype: _DTypeLike | None = None,
+        import_scope: str | None = None,
+        constraint: Callable[[Tensor], Tensor] | None = None,
+        synchronization: VariableSynchronization = VariableSynchronization.AUTO,
+        aggregation: VariableAggregation = VariableAggregation.NONE,
+        shape: _ShapeLike | None = None,
+    ): ...
+    def __getattr__(self, name: str) -> Incomplete: ...
+
+class RaggedTensor(metaclass=ABCMeta):
     def bounding_shape(
         self, axis: _TensorCompatible | None = None, name: str | None = None, out_type: _DTypeLike | None = None
     ) -> Tensor: ...
@@ -85,35 +127,42 @@ class RaggedTensor:
     def to_tensor(
         self, default_value: float | str | None = None, name: str | None = None, shape: _ShapeLike | None = None
     ) -> Tensor: ...
-    def __add__(self, other: RaggedTensor | float) -> RaggedTensor: ...
-    def __radd__(self, other: RaggedTensor | float) -> RaggedTensor: ...
-    def __sub__(self, other: RaggedTensor | float) -> RaggedTensor: ...
-    def __mul__(self, other: RaggedTensor | float) -> RaggedTensor: ...
-    def __rmul__(self, other: RaggedTensor | float) -> RaggedTensor: ...
-    def __floordiv__(self, other: RaggedTensor | float) -> RaggedTensor: ...
-    def __truediv__(self, other: RaggedTensor | float) -> RaggedTensor: ...
+    def __add__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
+    def __radd__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
+    def __sub__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
+    def __mul__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
+    def __rmul__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
+    def __floordiv__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
+    def __truediv__(self, other: RaggedTensor | float, name: str | None = None) -> RaggedTensor: ...
     def __getitem__(self, slice_spec: _Slice | tuple[_Slice, ...]) -> RaggedTensor: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class Operation:
+    def __init__(
+        self,
+        node_def: Incomplete,
+        g: Graph,
+        # isinstance is used so can not be Sequence/Iterable.
+        inputs: list[Tensor] | None = None,
+        output_types: Unused = None,
+        control_inputs: Iterable[Tensor | Operation] | None = None,
+        input_types: Iterable[DType] | None = None,
+        original_op: Operation | None = None,
+        op_def: Incomplete = None,
+    ): ...
     @property
     def inputs(self) -> list[Tensor]: ...
     @property
-    def input_types(self) -> list[DType]: ...
-    @property
     def outputs(self) -> list[Tensor]: ...
     @property
-    def output_types(self) -> list[DType]: ...
-    @property
     def device(self) -> str: ...
     @property
     def name(self) -> str: ...
     @property
     def type(self) -> str: ...
-    def __getitem__(self, slice_spec: int | slice | tuple[int | slice, ...]) -> Operation: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
-class TensorShape:
+class TensorShape(metaclass=ABCMeta):
     def __init__(self, dims: _ShapeLike) -> None: ...
     @property
     def rank(self) -> int: ...
diff --git a/stubs/tensorflow/tensorflow/core/framework/variable_pb2.pyi b/stubs/tensorflow/tensorflow/core/framework/variable_pb2.pyi
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/stubs/tensorflow/tensorflow/dtypes.pyi b/stubs/tensorflow/tensorflow/dtypes.pyi
index aec80b3be9ba..e119c1fece83 100644
--- a/stubs/tensorflow/tensorflow/dtypes.pyi
+++ b/stubs/tensorflow/tensorflow/dtypes.pyi
@@ -1,11 +1,14 @@
 from _typeshed import Incomplete
+from abc import ABCMeta
 from builtins import bool as _bool
 from typing import Any
 
 import numpy as np
 from tensorflow import _DTypeLike
 
-class DType:
+class _DTypeMeta(ABCMeta): ...
+
+class DType(metaclass=_DTypeMeta):
     @property
     def name(self) -> str: ...
     @property
@@ -48,5 +51,5 @@ quint8: DType = ...
 quint16: DType = ...
 string: DType = ...
 
-def as_dtype(dtype: _DTypeLike) -> DType: ...
+def as_dtype(type_value: _DTypeLike) -> DType: ...
 def __getattr__(name: str) -> Incomplete: ...
diff --git a/stubs/tensorflow/tensorflow/sparse.pyi b/stubs/tensorflow/tensorflow/sparse.pyi
index e1562644d003..625919c8b2f3 100644
--- a/stubs/tensorflow/tensorflow/sparse.pyi
+++ b/stubs/tensorflow/tensorflow/sparse.pyi
@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+from abc import ABCMeta
 from typing_extensions import TypeAlias
 
 from tensorflow import Tensor, TensorShape, _TensorCompatible
@@ -6,12 +7,17 @@ from tensorflow.dtypes import DType
 
 _SparseTensorCompatible: TypeAlias = _TensorCompatible | SparseTensor
 
-class SparseTensor:
-    indices: Tensor
-    values: Tensor
-    dense_shape: Tensor
-    shape: TensorShape
-    dtype: DType
+class SparseTensor(metaclass=ABCMeta):
+    @property
+    def indices(self) -> Tensor: ...
+    @property
+    def values(self) -> Tensor: ...
+    @property
+    def dense_shape(self) -> Tensor: ...
+    @property
+    def shape(self) -> TensorShape: ...
+    @property
+    def dtype(self) -> DType: ...
     name: str
     def __init__(self, indices: _TensorCompatible, values: _TensorCompatible, dense_shape: _TensorCompatible) -> None: ...
     def get_shape(self) -> TensorShape: ...
@@ -19,7 +25,6 @@ class SparseTensor:
     def __div__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
     def __truediv__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
     def __mul__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
-    def __rmul__(self, y: _SparseTensorCompatible) -> SparseTensor: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
 def __getattr__(name: str) -> Incomplete: ...

From 556fb08854a75c3bf6a96354cd7d6f70bc36f219 Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Tue, 10 Jan 2023 23:07:26 -0800
Subject: [PATCH 09/19] fix pyright

---
 stubs/tensorflow/tensorflow/__init__.pyi | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 2432959eae15..f798f2f5ada3 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -112,7 +112,7 @@ class Variable(Tensor, metaclass=_VariableMetaclass):
         synchronization: VariableSynchronization = VariableSynchronization.AUTO,
         aggregation: VariableAggregation = VariableAggregation.NONE,
         shape: _ShapeLike | None = None,
-    ): ...
+    ) -> None: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
 class RaggedTensor(metaclass=ABCMeta):
@@ -149,7 +149,7 @@ class Operation:
         input_types: Iterable[DType] | None = None,
         original_op: Operation | None = None,
         op_def: Incomplete = None,
-    ): ...
+    ) -> None: ...
     @property
     def inputs(self) -> list[Tensor]: ...
     @property

From 500ce70306e6d31445cacc63d0d00144da137b28 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 11 Jan 2023 11:42:51 +0000
Subject: [PATCH 10/19] Try to fix the mypy issue

---
 .github/workflows/tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index c36b24386d07..51f3391f3444 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -91,7 +91,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
-          python-version: "3.10"
+          python-version: ${{ matrix.python-version }}
           cache: pip
           cache-dependency-path: requirements-tests.txt
       - run: pip install -r requirements-tests.txt

From 1dd49f829012c3a21afc5b60f4c98f2da83689e5 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 11 Jan 2023 11:47:19 +0000
Subject: [PATCH 11/19] Update tests.yml

---
 .github/workflows/tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 51f3391f3444..1002f09374ca 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -94,7 +94,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
           cache: pip
           cache-dependency-path: requirements-tests.txt
-      - run: pip install -r requirements-tests.txt
+      - run: pip install mypy packaging pathspec termcolor tomli
       - run: python ./tests/mypy_test.py --platform=${{ matrix.platform }} --python-version=${{ matrix.python-version }}
 
   regression-tests:

From a20f2d9c4eb4a1a42e2b14d85dd6b87610af6846 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 11 Jan 2023 11:49:34 +0000
Subject: [PATCH 12/19] Update utils.py

---
 tests/utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/utils.py b/tests/utils.py
index a8d81f93e122..53c4b188f8c7 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -8,7 +8,7 @@
 import sys
 import venv
 from collections.abc import Iterable, Mapping
-from functools import cache
+from functools import lru_cache
 from pathlib import Path
 from typing import NamedTuple
 from typing_extensions import Annotated
@@ -24,6 +24,7 @@
     def colored(text: str, color: str | None = None, on_color: str | None = None, attrs: Iterable[str] | None = None) -> str:
         return text
 
+cache = lru_cache(None)
 
 # Used to install system-wide packages for different OS types:
 METADATA_MAPPING = {"linux": "apt_dependencies", "darwin": "brew_dependencies", "win32": "choco_dependencies"}

From 64bf119964150a8c460d65061a52c4d87b6d4755 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 11 Jan 2023 11:50:48 +0000
Subject: [PATCH 13/19] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 tests/utils.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/utils.py b/tests/utils.py
index 53c4b188f8c7..fdeaf253f871 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -24,6 +24,7 @@
     def colored(text: str, color: str | None = None, on_color: str | None = None, attrs: Iterable[str] | None = None) -> str:
         return text
 
+
 cache = lru_cache(None)
 
 # Used to install system-wide packages for different OS types:

From 4c6af9eb725ea0922a400e0ab3ab5f9ca75b8e20 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 11 Jan 2023 11:52:19 +0000
Subject: [PATCH 14/19] Update mypy_test.py

---
 tests/mypy_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/mypy_test.py b/tests/mypy_test.py
index 0fd07095b1d5..8e6a12d63d4c 100644
--- a/tests/mypy_test.py
+++ b/tests/mypy_test.py
@@ -16,7 +16,7 @@
 from itertools import product
 from pathlib import Path
 from threading import Lock
-from typing import TYPE_CHECKING, Any, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple, Tuple
 
 if TYPE_CHECKING:
     from _typeshed import StrPath
@@ -52,7 +52,7 @@
 
 ReturnCode: TypeAlias = int
 VersionString: TypeAlias = Annotated[str, "Must be one of the entries in SUPPORTED_VERSIONS"]
-VersionTuple: TypeAlias = tuple[int, int]
+VersionTuple: TypeAlias = Tuple[int, int]
 Platform: TypeAlias = Annotated[str, "Must be one of the entries in SUPPORTED_PLATFORMS"]
 
 

From 19ab96c9aec2bff1b6612528bce97c620e9497e5 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 11 Jan 2023 11:55:57 +0000
Subject: [PATCH 15/19] Update mypy_test.py

---
 tests/mypy_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/mypy_test.py b/tests/mypy_test.py
index 8e6a12d63d4c..2c23aa8a0195 100644
--- a/tests/mypy_test.py
+++ b/tests/mypy_test.py
@@ -323,7 +323,7 @@ def test_third_party_distribution(
 
     mypypath = os.pathsep.join(str(Path("stubs", dist)) for dist in seen_dists)
     if args.verbose:
-        print(colored(f"\n{mypypath=}", "blue"))
+        print(colored(f"\nmypypath={mypypath}", "blue"))
     code = run_mypy(
         args,
         configurations,

From db0924d49008ad6a98fc47af1cce2cd4ddb92875 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 11 Jan 2023 12:11:49 +0000
Subject: [PATCH 16/19] Update mypy_test.py

---
 tests/mypy_test.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/mypy_test.py b/tests/mypy_test.py
index 2c23aa8a0195..dafeb1d09245 100644
--- a/tests/mypy_test.py
+++ b/tests/mypy_test.py
@@ -12,6 +12,7 @@
 import tempfile
 import time
 from collections import defaultdict
+from collections.abc import Sequence
 from dataclasses import dataclass
 from itertools import product
 from pathlib import Path
@@ -77,6 +78,13 @@ def valid_path(cmd_arg: str) -> Path:
 parser = argparse.ArgumentParser(
     description="Typecheck typeshed's stubs with mypy. Patterns are unanchored regexps on the full path."
 )
+if sys.version_info < (3, 8):
+    class ExtendAction(argparse.Action):
+        def __call__(self, parser: argparse.ArgumentParser, namespace: argparse.Namespace, values: Sequence[str], option_string: object = None) -> None:
+            items = getattr(namespace, self.dest) or []
+            items.extend(values)
+            setattr(namespace, self.dest, items)
+    parser.register('action', 'extend', ExtendAction)
 parser.add_argument(
     "filter",
     type=valid_path,

From 9e01dde80a20f9514384a428b77b84656687352b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 11 Jan 2023 12:13:12 +0000
Subject: [PATCH 17/19] [pre-commit.ci] auto fixes from pre-commit.com hooks

---
 tests/mypy_test.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/tests/mypy_test.py b/tests/mypy_test.py
index dafeb1d09245..334a19529a3f 100644
--- a/tests/mypy_test.py
+++ b/tests/mypy_test.py
@@ -79,12 +79,20 @@ def valid_path(cmd_arg: str) -> Path:
     description="Typecheck typeshed's stubs with mypy. Patterns are unanchored regexps on the full path."
 )
 if sys.version_info < (3, 8):
+
     class ExtendAction(argparse.Action):
-        def __call__(self, parser: argparse.ArgumentParser, namespace: argparse.Namespace, values: Sequence[str], option_string: object = None) -> None:
+        def __call__(
+            self,
+            parser: argparse.ArgumentParser,
+            namespace: argparse.Namespace,
+            values: Sequence[str],
+            option_string: object = None,
+        ) -> None:
             items = getattr(namespace, self.dest) or []
             items.extend(values)
             setattr(namespace, self.dest, items)
-    parser.register('action', 'extend', ExtendAction)
+
+    parser.register("action", "extend", ExtendAction)
 parser.add_argument(
     "filter",
     type=valid_path,

From 150160d0a6bc6d65de0dbc870b96ad99a2c3cd6b Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@snapchat.com>
Date: Wed, 11 Jan 2023 09:08:24 -0800
Subject: [PATCH 18/19] Try to fix pytype with numpy usage.

---
 stubs/tensorflow/tensorflow/__init__.pyi | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index f798f2f5ada3..a3f284e0fe0c 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -7,7 +7,7 @@ from enum import Enum
 from typing import Any, NoReturn, overload
 from typing_extensions import TypeAlias
 
-import numpy as np
+import numpy
 from tensorflow.dtypes import *
 
 # Most tf.math functions are exported as tf, but sadly not all are.
@@ -23,10 +23,10 @@ _Slice: TypeAlias = int | slice | None
 
 _FloatDataSequence: TypeAlias = Sequence[float] | Sequence[_FloatDataSequence]
 _StrDataSequence: TypeAlias = Sequence[str] | Sequence[_StrDataSequence]
-_ScalarTensorCompatible: TypeAlias = Tensor | str | float | np.ndarray[Any, Any] | np.number[Any]
+_ScalarTensorCompatible: TypeAlias = Tensor | str | float | numpy.ndarray[Any, Any] | numpy.number[Any]
 _TensorCompatible: TypeAlias = _ScalarTensorCompatible | Sequence[_TensorCompatible]
 _ShapeLike: TypeAlias = TensorShape | Iterable[_ScalarTensorCompatible | None] | int | Tensor
-_DTypeLike: TypeAlias = DType | str | np.dtype[Any]
+_DTypeLike: TypeAlias = DType | str | numpy.dtype[Any]
 
 class Tensor:
     def __init__(self, op: Operation, value_index: int, dtype: DType) -> None: ...
@@ -42,7 +42,7 @@ class Tensor:
     def name(self) -> str: ...
     @property
     def op(self) -> Operation: ...
-    def numpy(self) -> np.ndarray[Any, Any]: ...
+    def numpy(self) -> numpy.ndarray[Any, Any]: ...
     def __int__(self) -> int: ...
     def __abs__(self, name: str | None = None) -> Tensor: ...
     def __add__(self, other: _TensorCompatible) -> Tensor: ...

From 5174b850b8f776295b39a8244caf90775c4ac796 Mon Sep 17 00:00:00 2001
From: Mehdi Drissi <mdrissi@hmc.edu>
Date: Wed, 11 Jan 2023 12:29:18 -0600
Subject: [PATCH 19/19] Apply suggestions from code review

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 stubs/tensorflow/METADATA.toml         |  3 +-
 stubs/tensorflow/tensorflow/dtypes.pyi | 46 +++++++++++++-------------
 2 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/stubs/tensorflow/METADATA.toml b/stubs/tensorflow/METADATA.toml
index 219877d7648b..396c33bb54ed 100644
--- a/stubs/tensorflow/METADATA.toml
+++ b/stubs/tensorflow/METADATA.toml
@@ -1,2 +1,3 @@
 version = "2.10.*"
-requires = ["numpy"]
\ No newline at end of file
+# requires a version of numpy with a `py.typed` file
+requires = ["numpy>=1.20"]
diff --git a/stubs/tensorflow/tensorflow/dtypes.pyi b/stubs/tensorflow/tensorflow/dtypes.pyi
index e119c1fece83..0ef3c1d295fa 100644
--- a/stubs/tensorflow/tensorflow/dtypes.pyi
+++ b/stubs/tensorflow/tensorflow/dtypes.pyi
@@ -27,29 +27,29 @@ class DType(metaclass=_DTypeMeta):
     def is_unsigned(self) -> _bool: ...
     def __getattr__(self, name: str) -> Incomplete: ...
 
-bool: DType = ...
-complex128: DType = ...
-complex64: DType = ...
-bfloat16: DType = ...
-float16: DType = ...
-half: DType = ...
-float32: DType = ...
-float64: DType = ...
-double: DType = ...
-int8: DType = ...
-int16: DType = ...
-int32: DType = ...
-int64: DType = ...
-uint8: DType = ...
-uint16: DType = ...
-uint32: DType = ...
-uint64: DType = ...
-qint8: DType = ...
-qint16: DType = ...
-qint32: DType = ...
-quint8: DType = ...
-quint16: DType = ...
-string: DType = ...
+bool: DType
+complex128: DType
+complex64: DType
+bfloat16: DType
+float16: DType
+half: DType
+float32: DType
+float64: DType
+double: DType
+int8: DType
+int16: DType
+int32: DType
+int64: DType
+uint8: DType
+uint16: DType
+uint32: DType
+uint64: DType
+qint8: DType
+qint16: DType
+qint32: DType
+quint8: DType
+quint16: DType
+string: DType
 
 def as_dtype(type_value: _DTypeLike) -> DType: ...
 def __getattr__(name: str) -> Incomplete: ...