Tensorflow keras layer (#9707)

Co-authored-by: Mehdi Drissi <mdrissi@snapchat.com>

Author: hmc-cs-mdrissi

stubs/tensorflow/@tests/stubtest_allowlist.txt

@@ -15,8 +15,52 @@ tensorflow.DType.__getattr__
 tensorflow.Graph.__getattr__
 tensorflow.Operation.__getattr__
 tensorflow.Variable.__getattr__
+tensorflow.keras.layers.Layer.__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__
+# stub internal utility
+tensorflow._aliases
+
+# Tensorflow imports are cursed.
+# import tensorflow.initializers
+# import tensorflow as tf
+# tf.initializers
+# Usually these two ways are same module, but for tensorflow the first way
+# often does not work and the second way does. The documentation describes
+# tf.initializers as module and has that type if accessed the second way,
+# but the real module file is completely different name (even package) and dynamically handled.
+# tf.initializers at runtime is <module 'keras.api._v2.keras.initializers' from '...'>
+tensorflow.initializers
+
+# Layer constructor's always have **kwargs, but only allow a few specific values. PEP 692
+# would allow us to specify this with **kwargs and remove the need for these exceptions.
+tensorflow.keras.layers.*.__init__
+
+# __call__ in tensorflow classes often allow keyword usage, but
+# when you subclass those classes it is not expected to handle keyword case. As an example,
+# class MyLayer(tf.keras.layers.Layer):
+#   def call(self, x):
+#      ...
+# is common even though Layer.call is defined like def call(self, inputs). Treating inputs as
+# a keyword argument would lead to many false positives with typical subclass usage.
+# Additional awkwardness for Layer's is call may optionally have training/mask as keyword arguments and some
+# layers do while others do not. At runtime call is not intended to be used directly by users,
+# but instead through __call__ which extracts out the training/mask arguments. Trying to describe
+# this better in stubs would similarly add many false positive Liskov violations.
+tensorflow.keras.layers.*.call
+tensorflow.keras.regularizers.Regularizer.__call__
+tensorflow.keras.constraints.Constraint.__call__
+
+# Layer class does good deal of __new__ magic and actually returns one of two different internal
+# types depending on tensorflow execution mode. This feels like implementation internal.
+tensorflow.keras.layers.Layer.__new__
+
+# build/compute_output_shape are marked positional only in stubs
+# as argument name is inconsistent across layer's and looks like
+# an implementation detail as documentation never mentions the
+# disagreements.
+tensorflow.keras.layers.*.build
+tensorflow.keras.layers.*.compute_output_shape

stubs/tensorflow/tensorflow/__init__.pyi

@@ -4,10 +4,12 @@ from builtins import bool as _bool
 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 Self, TypeAlias
+from types import TracebackType
+from typing import Any, NoReturn, TypeVar, overload
+from typing_extensions import ParamSpec, Self, TypeAlias
 
 import numpy
+from tensorflow import initializers as initializers, keras as keras, math as math
 
 # Explicit import of DType is covered by the wildcard, but
 # is necessary to avoid a crash in pytype.
@@ -253,4 +255,31 @@ class IndexedSlices(metaclass=ABCMeta):
     def __neg__(self) -> IndexedSlices: ...
     def consumers(self) -> list[Operation]: ...
 
+class name_scope:
+    def __init__(self, name: str) -> None: ...
+    def __enter__(self) -> str: ...
+    def __exit__(self, typ: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None) -> None: ...
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+class Module:
+    def __init__(self, name: str | None = None) -> None: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def name_scope(self) -> name_scope: ...
+    # Documentation only specifies these as returning Sequence. Actual
+    # implementation does tuple.
+    @property
+    def variables(self) -> Sequence[Variable]: ...
+    @property
+    def trainable_variables(self) -> Sequence[Variable]: ...
+    @property
+    def non_trainable_variables(self) -> Sequence[Variable]: ...
+    @property
+    def submodules(self) -> Sequence[Module]: ...
+    @classmethod
+    def with_name_scope(cls, method: Callable[_P, _R]) -> Callable[_P, _R]: ...
+
 def __getattr__(name: str) -> Incomplete: ...

stubs/tensorflow/tensorflow/_aliases.pyi

@@ -0,0 +1,14 @@
+# Commonly used type aliases.
+# Everything in this module is private for stubs. There is no runtime
+# equivalent.
+
+from collections.abc import Mapping, Sequence
+from typing import Any, TypeVar
+from typing_extensions import TypeAlias
+
+import numpy
+
+_T1 = TypeVar("_T1")
+ContainerGeneric: TypeAlias = Mapping[str, ContainerGeneric[_T1]] | Sequence[ContainerGeneric[_T1]] | _T1
+
+AnyArray: TypeAlias = numpy.ndarray[Any, Any]

stubs/tensorflow/tensorflow/initializers.pyi

@@ -0,0 +1 @@
+from tensorflow.keras.initializers import *

stubs/tensorflow/tensorflow/keras/__init__.pyi

@@ -0,0 +1,11 @@
+from _typeshed import Incomplete
+
+from tensorflow.keras import (
+    activations as activations,
+    constraints as constraints,
+    initializers as initializers,
+    layers as layers,
+    regularizers as regularizers,
+)
+
+def __getattr__(name: str) -> Incomplete: ...

stubs/tensorflow/tensorflow/keras/activations.pyi

@@ -0,0 +1,12 @@
+from _typeshed import Incomplete
+from collections.abc import Callable
+from typing import Any
+from typing_extensions import TypeAlias
+
+from tensorflow import Tensor
+
+# The implementation uses isinstance so it must be dict and not any Mapping.
+_Activation: TypeAlias = str | None | Callable[[Tensor], Tensor] | dict[str, Any]
+
+def get(identifier: _Activation) -> Callable[[Tensor], Tensor]: ...
+def __getattr__(name: str) -> Incomplete: ...

stubs/tensorflow/tensorflow/keras/constraints.pyi

@@ -0,0 +1,17 @@
+from _typeshed import Incomplete
+from collections.abc import Callable
+from typing import Any, overload
+
+from tensorflow import Tensor
+
+class Constraint:
+    def get_config(self) -> dict[str, Any]: ...
+    def __call__(self, __w: Tensor) -> Tensor: ...
+
+@overload
+def get(identifier: None) -> None: ...
+@overload
+def get(identifier: str | dict[str, Any] | Constraint) -> Constraint: ...
+@overload
+def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ...
+def __getattr__(name: str) -> Incomplete: ...

stubs/tensorflow/tensorflow/keras/initializers.pyi

@@ -0,0 +1,50 @@
+from _typeshed import Incomplete
+from collections.abc import Callable
+from typing import Any, overload
+from typing_extensions import Self, TypeAlias
+
+from tensorflow import Tensor, _DTypeLike, _ShapeLike, _TensorCompatible
+
+class Initializer:
+    def __call__(self, shape: _ShapeLike, dtype: _DTypeLike | None = None) -> Tensor: ...
+    def get_config(self) -> dict[str, Any]: ...
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> Self: ...
+
+class Constant(Initializer):
+    def __init__(self, value: _TensorCompatible = 0) -> None: ...
+
+class GlorotNormal(Initializer):
+    def __init__(self, seed: int | None = None) -> None: ...
+
+class GlorotUniform(Initializer):
+    def __init__(self, seed: int | None = None) -> None: ...
+
+class TruncatedNormal(Initializer):
+    def __init__(self, mean: _TensorCompatible = 0.0, stddev: _TensorCompatible = 0.05, seed: int | None = None) -> None: ...
+
+class RandomNormal(Initializer):
+    def __init__(self, mean: _TensorCompatible = 0.0, stddev: _TensorCompatible = 0.05, seed: int | None = None) -> None: ...
+
+class RandomUniform(Initializer):
+    def __init__(self, minval: _TensorCompatible = -0.05, maxval: _TensorCompatible = 0.05, seed: int | None = None) -> None: ...
+
+class Zeros(Initializer): ...
+
+constant = Constant
+glorot_normal = GlorotNormal
+glorot_uniform = GlorotUniform
+truncated_normal = TruncatedNormal
+zeros = Zeros
+
+_Initializer: TypeAlias = (  # noqa: Y047
+    str | Initializer | type[Initializer] | Callable[[_ShapeLike], Tensor] | dict[str, Any] | None
+)
+
+@overload
+def get(identifier: None) -> None: ...
+@overload
+def get(identifier: str | Initializer | dict[str, Any] | type[Initializer]) -> Initializer: ...
+@overload
+def get(identifier: Callable[[_ShapeLike], Tensor]) -> Callable[[_ShapeLike], Tensor]: ...
+def __getattr__(name: str) -> Incomplete: ...