From 71e7479ad7eb42ba29e22247d96970d2f555decd Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Fri, 10 Feb 2023 07:09:53 -0800 Subject: [PATCH 1/9] Add stubs for tensorflow Layer. --- stubs/tensorflow/tensorflow/_aliases.pyi | 14 ++ stubs/tensorflow/tensorflow/initializers.pyi | 1 + .../tensorflow/tensorflow/keras/__init__.pyi | 10 ++ .../tensorflow/keras/activations.pyi | 12 ++ .../tensorflow/keras/constraints.pyi | 22 +++ .../tensorflow/keras/initializers.pyi | 49 ++++++ stubs/tensorflow/tensorflow/keras/layers.pyi | 163 ++++++++++++++++++ .../tensorflow/keras/regularizers.pyi | 23 +++ 8 files changed, 294 insertions(+) create mode 100644 stubs/tensorflow/tensorflow/_aliases.pyi create mode 100644 stubs/tensorflow/tensorflow/initializers.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/__init__.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/activations.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/constraints.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/initializers.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/layers.pyi create mode 100644 stubs/tensorflow/tensorflow/keras/regularizers.pyi diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi new file mode 100644 index 000000000000..a15af94dea08 --- /dev/null +++ b/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] # noqa: Y047 diff --git a/stubs/tensorflow/tensorflow/initializers.pyi b/stubs/tensorflow/tensorflow/initializers.pyi new file mode 100644 index 000000000000..3caa7d18381d --- /dev/null +++ b/stubs/tensorflow/tensorflow/initializers.pyi @@ -0,0 +1 @@ +from tensorflow.keras.initializers import * diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi new file mode 100644 index 000000000000..43b1f2291811 --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi @@ -0,0 +1,10 @@ +from _typeshed import Incomplete + +from tensorflow.keras import ( + constraints as constraints, + initializers as initializers, + layers as layers, + regularizers as regularizers, +) + +def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/activations.pyi b/stubs/tensorflow/tensorflow/keras/activations.pyi new file mode 100644 index 000000000000..1e90deb64642 --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/activations.pyi @@ -0,0 +1,12 @@ +from _typeshed import Incomplete +from collections.abc import Callable +from typing import Any, TypedDict +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] | TypedDict + +def get(identifier: _Activation) -> Callable[[Tensor], Tensor]: ... +def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/constraints.pyi b/stubs/tensorflow/tensorflow/keras/constraints.pyi new file mode 100644 index 000000000000..6a31692f774f --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/constraints.pyi @@ -0,0 +1,22 @@ +from _typeshed import Incomplete +from collections.abc import Callable +from typing import Any, overload +from typing_extensions import TypeAlias + +from tensorflow import Tensor + +class Constraint: + def get_config(self) -> dict[str, Any]: ... + def __call__(self, w: Tensor, /) -> Tensor: ... + +_Constraint: TypeAlias = str | dict[str, Any] | Constraint | None + +@overload +def get(identifer: None) -> None: ... +@overload +def get(identifer: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... +@overload +def get(identifer: str | dict[str, Any]) -> Constraint: ... +@overload +def get(identifer: _Constraint) -> Constraint | Callable[[Tensor], Tensor] | None: ... +def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/initializers.pyi b/stubs/tensorflow/tensorflow/keras/initializers.pyi new file mode 100644 index 000000000000..8d3b8487238f --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/initializers.pyi @@ -0,0 +1,49 @@ +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: type[Self], config: dict[str, Any]) -> Self: ... + +class Constant(Initializer): + def __init__(self, value: _TensorCompatible) -> 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): ... + +class RandomNormal(Initializer): + def __init__(self, mean: _TensorCompatible = 0.0, stddev: _TensorCompatible = 0.05, seed: int | None = None): ... + +class RandomUniform(Initializer): + def __init__(self, minval: _TensorCompatible = 0.0, maxval: _TensorCompatible = 0.05, seed: int | None = None): ... + +class Zeros(Initializer): + def __init__(self) -> None: ... + +constant = Constant +glorot_normal = GlorotNormal +glorot_uniform = GlorotUniform +truncated_normal = TruncatedNormal +zeros = Zeros + +_Initializer: TypeAlias = str | Initializer | Callable[[_ShapeLike], Tensor] | dict[str, Any] | None + +@overload +def get(identifier: None) -> None: ... +@overload +def get(identifier: str | Initializer | dict[str, Any]) -> Initializer: ... +@overload +def get(identifier: _Initializer) -> Initializer | None: ... +def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi new file mode 100644 index 000000000000..56badf4dec41 --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -0,0 +1,163 @@ +from _typeshed import Incomplete +from collections.abc import Callable, Iterable, Sequence +from typing import Any, Generic, TypeVar, overload +from typing_extensions import Self, TypeAlias + +import tensorflow as tf +from tensorflow import DTypeLike, Tensor, Variable, _TensorCompatible +from tensorflow._aliases import _AnyArray, _ContainerGeneric +from tensorflow.keras.activations import _Activation +from tensorflow.keras.constraints import _Constraint +from tensorflow.keras.initializers import _Initializer +from tensorflow.keras.regularizers import _Regularizer + +_InputT = TypeVar("_InputT", contravariant=True) +_OutputT = TypeVar("_OutputT", covariant=True) + +class InputSpec: + dtype: str | None + shape: tuple[int | None, ...] + ndim: int | None + max_ndim: int | None + min_ndim: int | None + axes: dict[int, int | None] | None + def __init__( + self, + dtype: tf._DTypeLike | None = None, + shape: Iterable[int | None] | None = None, + ndim: int | None = None, + max_ndim: int | None = None, + min_ndim: int | None = None, + axes: dict[int, int | None] | None = None, + allow_last_axis_squeeze: bool = False, + name: str | None = None, + ): ... + def get_config(self) -> dict[str, Any]: ... + @classmethod + def from_config(cls, config: dict[str, Any]) -> type[Self]: ... + +_InputSpecs: TypeAlias = _ContainerGeneric[InputSpec] + +# Most layers have input and output type of just Tensor and when we support default type variables, +# maybe worth trying. +class Layer(Generic[_InputT, _OutputT], tf.Module): + input_spec: _InputSpecs + + @property + def trainable(self) -> bool: ... + @trainable.setter + def trainable(self, value: bool) -> None: ... + def __init__( + self, trainable: bool = True, name: str | None = None, dtype: DTypeLike | None = None, dynamic: bool = False + ) -> None: ... + + # *args/**kwargs are allowed, but have obscure footguns and tensorflow documentation discourages their usage. + # First argument will automatically be cast to layer's compute dtype, but any other tensor arguments will not be. + # Also various tensorflow tools/apis can misbehave if they encounter a layer with *args/**kwargs. + def __call__(self, inputs: _InputT, /, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ... + def call(self, inputs: _InputT, /) -> _OutputT: ... + + # input_shape's real type depends on _InputT, but we can't express that without HKT. + # For example _InputT tf.Tensor -> tf.TensorShape, _InputT dict[str, tf.Tensor] -> dict[str, tf.TensorShape]. + def build(self, input_shape: Any) -> None: ... + @overload + def compute_output_shape(self: Layer[tf.Tensor, tf.Tensor], input_shape: tf.TensorShape) -> tf.TensorShape: ... + @overload + def compute_output_shape(self, input_shape: Any) -> Any: ... + def add_weight( + self, + name: str | None = None, + shape: Iterable[int | None] | None = None, + dtype: tf._DTypeLike | None = None, + initializer: _Initializer | None = None, + regularizer: _Regularizer = None, + constraint: _Constraint = None, + trainable: bool | None = None, + ) -> tf.Variable: ... + def add_loss(self, losses: tf.Tensor | Sequence[tf.Tensor] | Callable[[], tf.Tensor]) -> None: ... + def count_params(self) -> int: ... + @property + def trainable_variables(self) -> list[Variable]: ... + @property + def non_trainable_variables(self) -> list[Variable]: ... + @property + def trainable_weights(self) -> list[Variable]: ... + @property + def non_trainable_weights(self) -> list[Variable]: ... + @property + def losses(self) -> list[Tensor]: ... + def get_weights(self) -> list[_AnyArray]: ... + def set_weights(self, weights: Sequence[_AnyArray]) -> None: ... + def get_config(self) -> dict[str, Any]: ... + @classmethod + def from_config(cls: type[Self], config: dict[str, Any]) -> Self: ... + def __getattr__(name: str) -> Incomplete: ... + +class Dense(Layer[tf.Tensor, tf.Tensor]): + input_spec: InputSpec + def __init__( + self, + units: int, + activation: _Activation = None, + use_bias: bool = True, + kernel_initializer: _Initializer = "glorot_uniform", + bias_initializer: _Initializer = "zeros", + kernel_regularizer: _Regularizer = None, + bias_regularizer: _Regularizer = None, + activity_regularizer: _Regularizer = None, + kernel_constraint: _Constraint = None, + bias_constraint: _Constraint = None, + name: str | None = None, + ) -> None: ... + +class BatchNormalization(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, + axis: int = -1, + momentum: float = 0.99, + epsilon: float = 0.001, + center: bool = True, + scale: bool = True, + beta_initializer: _Initializer = "zeros", + gamma_initializer: _Initializer = "ones", + moving_mean_initializer: _Initializer = "zeros", + moving_variance_initializer: _Initializer = "ones", + beta_regularizer: _Regularizer = None, + gamma_regularizer: _Regularizer = None, + beta_constraint: _Constraint = None, + gamma_constraint: _Constraint = None, + name: str | None = None, + ): ... + +class ReLU(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, + max_value: float | None = None, + negative_slope: float | None = 0.0, + threshold: float | None = 0.0, + name: str | None = None, + ) -> None: ... + +class Dropout(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, + rate: float, + noise_shape: _TensorCompatible | Sequence[int | None] | None = None, + seed: int | None = None, + name: str | None = None, + ) -> None: ... + +class Embedding(Layer[tf.Tensor, tf.Tensor]): + def __init__( + self, + input_dim: int, + output_dim: int, + embeddings_initializer: _Initializer = "uniform", + embeddings_regularizer: _Regularizer = None, + embeddings_constraint: _Constraint = None, + mask_zero: bool = False, + input_length: int | None = None, + name: str | None = None, + ): ... + +def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/regularizers.pyi b/stubs/tensorflow/tensorflow/keras/regularizers.pyi new file mode 100644 index 000000000000..cffad88b2d06 --- /dev/null +++ b/stubs/tensorflow/tensorflow/keras/regularizers.pyi @@ -0,0 +1,23 @@ +from collections.abc import Callable +from typing import Any, overload +from typing_extensions import Self, TypeAlias + +from tensorflow import Tensor + +class Regularizer: + def get_config(self) -> dict[str, Any]: ... + @classmethod + def from_config(cls, config: dict[str, Any]) -> Self: ... + def __call__(self, x: Tensor, /) -> Tensor: ... + +_Regularizer: TypeAlias = str | dict[str, Any] | Regularizer | None + +@overload +def get(identifer: None) -> None: ... +@overload +def get(identifer: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... +@overload +def get(identifer: str | dict[str, Any]) -> Regularizer: ... +@overload +def get(identifer: _Regularizer) -> Regularizer | Callable[[Tensor], Tensor] | None: ... +def __getattr__(name: str) -> Any: ... From deb5c7dfa7f4d78cdeb857d9fdc2070a2cfb14db Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Fri, 10 Feb 2023 08:51:26 -0800 Subject: [PATCH 2/9] Get stubtest to pass for tf layer. --- .../tensorflow/@tests/stubtest_allowlist.txt | 41 ++++++++++++++++ stubs/tensorflow/tensorflow/__init__.pyi | 32 ++++++++++++- .../tensorflow/keras/activations.pyi | 4 +- .../tensorflow/keras/constraints.pyi | 10 ++-- .../tensorflow/keras/initializers.pyi | 9 ++-- stubs/tensorflow/tensorflow/keras/layers.pyi | 47 ++++++++++++++++--- .../tensorflow/keras/regularizers.pyi | 10 ++-- 7 files changed, 128 insertions(+), 25 deletions(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index e42fc68f1b4c..657bc184bd86 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt @@ -15,8 +15,49 @@ 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 +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.Dense.__init__ +tensorflow.keras.layers.BatchNormalization.__init__ +tensorflow.keras.layers.ReLU.__init__ +tensorflow.keras.layers.Dropout.__init__ +tensorflow.keras.layers.Embedding.__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__ diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 62cab55b71a2..5e4031b0e11f 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -4,8 +4,9 @@ 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 TypeAlias +from types import TracebackType +from typing import Any, NoReturn, TypeVar, overload +from typing_extensions import ParamSpec, TypeAlias import numpy @@ -253,4 +254,31 @@ class IndexedSlices(metaclass=ABCMeta): def __neg__(self) -> IndexedSlices: ... def consumers(self) -> list[Operation]: ... +class name_scope: + def __init__(self, name: str): ... + 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: ... diff --git a/stubs/tensorflow/tensorflow/keras/activations.pyi b/stubs/tensorflow/tensorflow/keras/activations.pyi index 1e90deb64642..650a449ae898 100644 --- a/stubs/tensorflow/tensorflow/keras/activations.pyi +++ b/stubs/tensorflow/tensorflow/keras/activations.pyi @@ -1,12 +1,12 @@ from _typeshed import Incomplete from collections.abc import Callable -from typing import Any, TypedDict +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] | TypedDict +_Activation: TypeAlias = str | None | Callable[[Tensor], Tensor] | dict[str, Any] def get(identifier: _Activation) -> Callable[[Tensor], Tensor]: ... def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/constraints.pyi b/stubs/tensorflow/tensorflow/keras/constraints.pyi index 6a31692f774f..a7fd7d3adc9c 100644 --- a/stubs/tensorflow/tensorflow/keras/constraints.pyi +++ b/stubs/tensorflow/tensorflow/keras/constraints.pyi @@ -7,16 +7,16 @@ from tensorflow import Tensor class Constraint: def get_config(self) -> dict[str, Any]: ... - def __call__(self, w: Tensor, /) -> Tensor: ... + def __call__(self, __w: Tensor) -> Tensor: ... _Constraint: TypeAlias = str | dict[str, Any] | Constraint | None @overload -def get(identifer: None) -> None: ... +def get(identifier: None) -> None: ... @overload -def get(identifer: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... +def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... @overload -def get(identifer: str | dict[str, Any]) -> Constraint: ... +def get(identifier: str | dict[str, Any]) -> Constraint: ... @overload -def get(identifer: _Constraint) -> Constraint | Callable[[Tensor], Tensor] | None: ... +def get(identifier: _Constraint) -> Constraint | Callable[[Tensor], Tensor] | None: ... def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/initializers.pyi b/stubs/tensorflow/tensorflow/keras/initializers.pyi index 8d3b8487238f..e4250fe5bb3a 100644 --- a/stubs/tensorflow/tensorflow/keras/initializers.pyi +++ b/stubs/tensorflow/tensorflow/keras/initializers.pyi @@ -9,10 +9,10 @@ class Initializer: def __call__(self, shape: _ShapeLike, dtype: _DTypeLike | None = None) -> Tensor: ... def get_config(self) -> dict[str, Any]: ... @classmethod - def from_config(cls: type[Self], config: dict[str, Any]) -> Self: ... + def from_config(cls, config: dict[str, Any]) -> Self: ... class Constant(Initializer): - def __init__(self, value: _TensorCompatible) -> None: ... + def __init__(self, value: _TensorCompatible = 0) -> None: ... class GlorotNormal(Initializer): def __init__(self, seed: int | None = None) -> None: ... @@ -27,10 +27,9 @@ class RandomNormal(Initializer): def __init__(self, mean: _TensorCompatible = 0.0, stddev: _TensorCompatible = 0.05, seed: int | None = None): ... class RandomUniform(Initializer): - def __init__(self, minval: _TensorCompatible = 0.0, maxval: _TensorCompatible = 0.05, seed: int | None = None): ... + def __init__(self, minval: _TensorCompatible = -0.05, maxval: _TensorCompatible = 0.05, seed: int | None = None): ... -class Zeros(Initializer): - def __init__(self) -> None: ... +class Zeros(Initializer): ... constant = Constant glorot_normal = GlorotNormal diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi index 56badf4dec41..36dc4b23fe44 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -4,7 +4,7 @@ from typing import Any, Generic, TypeVar, overload from typing_extensions import Self, TypeAlias import tensorflow as tf -from tensorflow import DTypeLike, Tensor, Variable, _TensorCompatible +from tensorflow import DTypeLike, Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible from tensorflow._aliases import _AnyArray, _ContainerGeneric from tensorflow.keras.activations import _Activation from tensorflow.keras.constraints import _Constraint @@ -54,8 +54,8 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): # *args/**kwargs are allowed, but have obscure footguns and tensorflow documentation discourages their usage. # First argument will automatically be cast to layer's compute dtype, but any other tensor arguments will not be. # Also various tensorflow tools/apis can misbehave if they encounter a layer with *args/**kwargs. - def __call__(self, inputs: _InputT, /, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ... - def call(self, inputs: _InputT, /) -> _OutputT: ... + def __call__(self, inputs: _InputT, *, training: bool = False, mask: _TensorCompatible | None = None) -> _OutputT: ... + def call(self, __inputs: _InputT) -> _OutputT: ... # input_shape's real type depends on _InputT, but we can't express that without HKT. # For example _InputT tf.Tensor -> tf.TensorShape, _InputT dict[str, tf.Tensor] -> dict[str, tf.TensorShape]. @@ -71,8 +71,11 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): dtype: tf._DTypeLike | None = None, initializer: _Initializer | None = None, regularizer: _Regularizer = None, - constraint: _Constraint = None, trainable: bool | None = None, + constraint: _Constraint = None, + use_resource: bool | None = None, + synchronization: VariableSynchronization = ..., + aggregation: VariableAggregation = ..., ) -> tf.Variable: ... def add_loss(self, losses: tf.Tensor | Sequence[tf.Tensor] | Callable[[], tf.Tensor]) -> None: ... def count_params(self) -> int: ... @@ -90,9 +93,23 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): def set_weights(self, weights: Sequence[_AnyArray]) -> None: ... def get_config(self) -> dict[str, Any]: ... @classmethod - def from_config(cls: type[Self], config: dict[str, Any]) -> Self: ... - def __getattr__(name: str) -> Incomplete: ... + def from_config(cls, config: dict[str, Any]) -> Self: ... + def __getattr__(self, name: str) -> Incomplete: ... + +# Every layer has trainable, dtype, name, and dynamic. At runtime these +# are mainly handled with **kwargs, passed up and then validated. +# In actual implementation there's 12 allowed keyword arguments, but only +# 4 are documented and other 8 are mainly internal. The other 8 can be found +# https://github.com/keras-team/keras/blob/e6784e4302c7b8cd116b74a784f4b78d60e83c26/keras/engine/base_layer.py#L329 +# PEP 692 support would be very helpful here and allow removing stubtest allowlist for +# all layer constructors. + +# TODO: Replace last Any after adding tf.keras.mixed_precision.Policy. +_LayerDtype = tf._DTypeLike | dict[str, Any] | Any +# Layer's compute_output_shape commonly have instance as first argument name instead of self. +# This is an artifact of actual implementation commonly uses a decorator to define it. +# Layer.build has same weirdness sometimes. class Dense(Layer[tf.Tensor, tf.Tensor]): input_spec: InputSpec def __init__( @@ -107,6 +124,9 @@ class Dense(Layer[tf.Tensor, tf.Tensor]): activity_regularizer: _Regularizer = None, kernel_constraint: _Constraint = None, bias_constraint: _Constraint = None, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -126,6 +146,9 @@ class BatchNormalization(Layer[tf.Tensor, tf.Tensor]): gamma_regularizer: _Regularizer = None, beta_constraint: _Constraint = None, gamma_constraint: _Constraint = None, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, name: str | None = None, ): ... @@ -135,8 +158,12 @@ class ReLU(Layer[tf.Tensor, tf.Tensor]): max_value: float | None = None, negative_slope: float | None = 0.0, threshold: float | None = 0.0, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, name: str | None = None, ) -> None: ... + def compute_output_shape(instance, input_shape: tf.TensorShape) -> tf.TensorShape: ... class Dropout(Layer[tf.Tensor, tf.Tensor]): def __init__( @@ -144,6 +171,9 @@ class Dropout(Layer[tf.Tensor, tf.Tensor]): rate: float, noise_shape: _TensorCompatible | Sequence[int | None] | None = None, seed: int | None = None, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, name: str | None = None, ) -> None: ... @@ -157,7 +187,12 @@ class Embedding(Layer[tf.Tensor, tf.Tensor]): embeddings_constraint: _Constraint = None, mask_zero: bool = False, input_length: int | None = None, + trainable: bool = True, + dtype: _LayerDtype = None, + dynamic: bool = False, name: str | None = None, ): ... + def build(instance, input_shape: tf.TensorShape) -> None: ... + def compute_output_shape(instance, input_shape: tf.TensorShape) -> tf.TensorShape: ... def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/regularizers.pyi b/stubs/tensorflow/tensorflow/keras/regularizers.pyi index cffad88b2d06..bc00770db765 100644 --- a/stubs/tensorflow/tensorflow/keras/regularizers.pyi +++ b/stubs/tensorflow/tensorflow/keras/regularizers.pyi @@ -8,16 +8,16 @@ class Regularizer: def get_config(self) -> dict[str, Any]: ... @classmethod def from_config(cls, config: dict[str, Any]) -> Self: ... - def __call__(self, x: Tensor, /) -> Tensor: ... + def __call__(self, __x: Tensor) -> Tensor: ... _Regularizer: TypeAlias = str | dict[str, Any] | Regularizer | None @overload -def get(identifer: None) -> None: ... +def get(identifier: None) -> None: ... @overload -def get(identifer: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... +def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... @overload -def get(identifer: str | dict[str, Any]) -> Regularizer: ... +def get(identifier: str | dict[str, Any]) -> Regularizer: ... @overload -def get(identifer: _Regularizer) -> Regularizer | Callable[[Tensor], Tensor] | None: ... +def get(identifier: _Regularizer) -> Regularizer | Callable[[Tensor], Tensor] | None: ... def __getattr__(name: str) -> Any: ... From 0e76cfd4f129d23db1d01f8624cd6e7642f64707 Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Fri, 10 Feb 2023 09:15:01 -0800 Subject: [PATCH 3/9] fix ci errors --- stubs/tensorflow/@tests/stubtest_allowlist.txt | 6 +----- stubs/tensorflow/tensorflow/__init__.pyi | 2 +- stubs/tensorflow/tensorflow/keras/initializers.pyi | 6 +++--- stubs/tensorflow/tensorflow/keras/layers.pyi | 12 ++++++------ 4 files changed, 11 insertions(+), 15 deletions(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index 657bc184bd86..febfb497cfd2 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt @@ -37,11 +37,7 @@ 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.Dense.__init__ -tensorflow.keras.layers.BatchNormalization.__init__ -tensorflow.keras.layers.ReLU.__init__ -tensorflow.keras.layers.Dropout.__init__ -tensorflow.keras.layers.Embedding.__init__ +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, diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 5e4031b0e11f..818898bb70cd 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -255,7 +255,7 @@ class IndexedSlices(metaclass=ABCMeta): def consumers(self) -> list[Operation]: ... class name_scope: - def __init__(self, name: str): ... + def __init__(self, name: str) -> None: ... def __enter__(self) -> str: ... def __exit__(self, typ: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None) -> None: ... diff --git a/stubs/tensorflow/tensorflow/keras/initializers.pyi b/stubs/tensorflow/tensorflow/keras/initializers.pyi index e4250fe5bb3a..5f23acf00b30 100644 --- a/stubs/tensorflow/tensorflow/keras/initializers.pyi +++ b/stubs/tensorflow/tensorflow/keras/initializers.pyi @@ -21,13 +21,13 @@ 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): ... + 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): ... + 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): ... + def __init__(self, minval: _TensorCompatible = -0.05, maxval: _TensorCompatible = 0.05, seed: int | None = None) -> None: ... class Zeros(Initializer): ... diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi index 36dc4b23fe44..0bcb4baab9fd 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -4,7 +4,7 @@ from typing import Any, Generic, TypeVar, overload from typing_extensions import Self, TypeAlias import tensorflow as tf -from tensorflow import DTypeLike, Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible +from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible from tensorflow._aliases import _AnyArray, _ContainerGeneric from tensorflow.keras.activations import _Activation from tensorflow.keras.constraints import _Constraint @@ -31,7 +31,7 @@ class InputSpec: axes: dict[int, int | None] | None = None, allow_last_axis_squeeze: bool = False, name: str | None = None, - ): ... + ) -> None: ... def get_config(self) -> dict[str, Any]: ... @classmethod def from_config(cls, config: dict[str, Any]) -> type[Self]: ... @@ -48,7 +48,7 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): @trainable.setter def trainable(self, value: bool) -> None: ... def __init__( - self, trainable: bool = True, name: str | None = None, dtype: DTypeLike | None = None, dynamic: bool = False + self, trainable: bool = True, name: str | None = None, dtype: tf._DTypeLike | None = None, dynamic: bool = False ) -> None: ... # *args/**kwargs are allowed, but have obscure footguns and tensorflow documentation discourages their usage. @@ -105,7 +105,7 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): # all layer constructors. # TODO: Replace last Any after adding tf.keras.mixed_precision.Policy. -_LayerDtype = tf._DTypeLike | dict[str, Any] | Any +_LayerDtype: TypeAlias = tf._DTypeLike | dict[str, Any] | Any # Layer's compute_output_shape commonly have instance as first argument name instead of self. # This is an artifact of actual implementation commonly uses a decorator to define it. @@ -150,7 +150,7 @@ class BatchNormalization(Layer[tf.Tensor, tf.Tensor]): dtype: _LayerDtype = None, dynamic: bool = False, name: str | None = None, - ): ... + ) -> None: ... class ReLU(Layer[tf.Tensor, tf.Tensor]): def __init__( @@ -191,7 +191,7 @@ class Embedding(Layer[tf.Tensor, tf.Tensor]): dtype: _LayerDtype = None, dynamic: bool = False, name: str | None = None, - ): ... + ) -> None: ... def build(instance, input_shape: tf.TensorShape) -> None: ... def compute_output_shape(instance, input_shape: tf.TensorShape) -> tf.TensorShape: ... From 7bcb56c2d85f9d4ff169173e6f7b5c2c9873b6da Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Sat, 11 Feb 2023 07:07:22 -0800 Subject: [PATCH 4/9] Add a couple re-exports --- stubs/tensorflow/tensorflow/__init__.pyi | 1 + stubs/tensorflow/tensorflow/keras/__init__.pyi | 1 + 2 files changed, 2 insertions(+) diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi index 818898bb70cd..fbfc4709a7d1 100644 --- a/stubs/tensorflow/tensorflow/__init__.pyi +++ b/stubs/tensorflow/tensorflow/__init__.pyi @@ -9,6 +9,7 @@ from typing import Any, NoReturn, TypeVar, overload from typing_extensions import ParamSpec, 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. diff --git a/stubs/tensorflow/tensorflow/keras/__init__.pyi b/stubs/tensorflow/tensorflow/keras/__init__.pyi index 43b1f2291811..c96602bc2c47 100644 --- a/stubs/tensorflow/tensorflow/keras/__init__.pyi +++ b/stubs/tensorflow/tensorflow/keras/__init__.pyi @@ -1,6 +1,7 @@ from _typeshed import Incomplete from tensorflow.keras import ( + activations as activations, constraints as constraints, initializers as initializers, layers as layers, From 6a241d1acfef0209fb19fff253a0e7834fe8d827 Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Sun, 12 Feb 2023 04:52:50 -0800 Subject: [PATCH 5/9] Relax input spec type. --- stubs/tensorflow/tensorflow/keras/layers.pyi | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi index 0bcb4baab9fd..d56f9e3b3812 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -5,7 +5,7 @@ from typing_extensions import Self, TypeAlias import tensorflow as tf from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible -from tensorflow._aliases import _AnyArray, _ContainerGeneric +from tensorflow._aliases import _AnyArray from tensorflow.keras.activations import _Activation from tensorflow.keras.constraints import _Constraint from tensorflow.keras.initializers import _Initializer @@ -36,12 +36,13 @@ class InputSpec: @classmethod def from_config(cls, config: dict[str, Any]) -> type[Self]: ... -_InputSpecs: TypeAlias = _ContainerGeneric[InputSpec] - # Most layers have input and output type of just Tensor and when we support default type variables, # maybe worth trying. class Layer(Generic[_InputT, _OutputT], tf.Module): - input_spec: _InputSpecs + # The most general type is _ContainerGeneric[InputSpec] as it really + # depends on _InputT. For most Layers it is just InputSpec + # though. Maybe describable with HKT? + input_spec: InputSpec | Any @property def trainable(self) -> bool: ... @@ -111,7 +112,6 @@ _LayerDtype: TypeAlias = tf._DTypeLike | dict[str, Any] | Any # This is an artifact of actual implementation commonly uses a decorator to define it. # Layer.build has same weirdness sometimes. class Dense(Layer[tf.Tensor, tf.Tensor]): - input_spec: InputSpec def __init__( self, units: int, From 06e9bb219246328dd37a63e63a37b953b906026c Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Wed, 22 Feb 2023 00:19:07 -0800 Subject: [PATCH 6/9] Review feedback. --- stubs/tensorflow/tensorflow/_aliases.pyi | 4 ++-- stubs/tensorflow/tensorflow/keras/constraints.pyi | 6 ++---- stubs/tensorflow/tensorflow/keras/initializers.pyi | 8 +++++--- stubs/tensorflow/tensorflow/keras/layers.pyi | 6 +++--- stubs/tensorflow/tensorflow/keras/regularizers.pyi | 6 ++---- 5 files changed, 14 insertions(+), 16 deletions(-) diff --git a/stubs/tensorflow/tensorflow/_aliases.pyi b/stubs/tensorflow/tensorflow/_aliases.pyi index a15af94dea08..32fe58c170f5 100644 --- a/stubs/tensorflow/tensorflow/_aliases.pyi +++ b/stubs/tensorflow/tensorflow/_aliases.pyi @@ -9,6 +9,6 @@ from typing_extensions import TypeAlias import numpy _T1 = TypeVar("_T1") -_ContainerGeneric: TypeAlias = Mapping[str, _ContainerGeneric[_T1]] | Sequence[_ContainerGeneric[_T1]] | _T1 +ContainerGeneric: TypeAlias = Mapping[str, ContainerGeneric[_T1]] | Sequence[ContainerGeneric[_T1]] | _T1 -_AnyArray: TypeAlias = numpy.ndarray[Any, Any] # noqa: Y047 +AnyArray: TypeAlias = numpy.ndarray[Any, Any] diff --git a/stubs/tensorflow/tensorflow/keras/constraints.pyi b/stubs/tensorflow/tensorflow/keras/constraints.pyi index a7fd7d3adc9c..4ef962cdf835 100644 --- a/stubs/tensorflow/tensorflow/keras/constraints.pyi +++ b/stubs/tensorflow/tensorflow/keras/constraints.pyi @@ -9,14 +9,12 @@ class Constraint: def get_config(self) -> dict[str, Any]: ... def __call__(self, __w: Tensor) -> Tensor: ... -_Constraint: TypeAlias = str | dict[str, Any] | Constraint | None +_Constraint: TypeAlias = str | dict[str, Any] | Constraint | None # noqa: Y047 @overload def get(identifier: None) -> None: ... @overload def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... @overload -def get(identifier: str | dict[str, Any]) -> Constraint: ... -@overload -def get(identifier: _Constraint) -> Constraint | Callable[[Tensor], Tensor] | None: ... +def get(identifier: str | dict[str, Any] | Constraint) -> Constraint: ... def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/initializers.pyi b/stubs/tensorflow/tensorflow/keras/initializers.pyi index 5f23acf00b30..f4aef5ce0002 100644 --- a/stubs/tensorflow/tensorflow/keras/initializers.pyi +++ b/stubs/tensorflow/tensorflow/keras/initializers.pyi @@ -37,12 +37,14 @@ glorot_uniform = GlorotUniform truncated_normal = TruncatedNormal zeros = Zeros -_Initializer: TypeAlias = str | Initializer | Callable[[_ShapeLike], Tensor] | dict[str, Any] | None +_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]) -> Initializer: ... +def get(identifier: str | Initializer | dict[str, Any] | type[Initializer]) -> Initializer: ... @overload -def get(identifier: _Initializer) -> Initializer | None: ... +def get(identifier: Callable[[_ShapeLike], Tensor]) -> Callable[[_ShapeLike], Tensor]: ... def __getattr__(name: str) -> Incomplete: ... diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi index d56f9e3b3812..993a3298e6e9 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -5,7 +5,7 @@ from typing_extensions import Self, TypeAlias import tensorflow as tf from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible -from tensorflow._aliases import _AnyArray +from tensorflow._aliases import AnyArray from tensorflow.keras.activations import _Activation from tensorflow.keras.constraints import _Constraint from tensorflow.keras.initializers import _Initializer @@ -90,8 +90,8 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): def non_trainable_weights(self) -> list[Variable]: ... @property def losses(self) -> list[Tensor]: ... - def get_weights(self) -> list[_AnyArray]: ... - def set_weights(self, weights: Sequence[_AnyArray]) -> None: ... + def get_weights(self) -> list[AnyArray]: ... + def set_weights(self, weights: Sequence[AnyArray]) -> None: ... def get_config(self) -> dict[str, Any]: ... @classmethod def from_config(cls, config: dict[str, Any]) -> Self: ... diff --git a/stubs/tensorflow/tensorflow/keras/regularizers.pyi b/stubs/tensorflow/tensorflow/keras/regularizers.pyi index bc00770db765..e73df6b0eda1 100644 --- a/stubs/tensorflow/tensorflow/keras/regularizers.pyi +++ b/stubs/tensorflow/tensorflow/keras/regularizers.pyi @@ -10,14 +10,12 @@ class Regularizer: def from_config(cls, config: dict[str, Any]) -> Self: ... def __call__(self, __x: Tensor) -> Tensor: ... -_Regularizer: TypeAlias = str | dict[str, Any] | Regularizer | None +_Regularizer: TypeAlias = str | dict[str, Any] | Regularizer | None # noqa: Y047 @overload def get(identifier: None) -> None: ... @overload def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... @overload -def get(identifier: str | dict[str, Any]) -> Regularizer: ... -@overload -def get(identifier: _Regularizer) -> Regularizer | Callable[[Tensor], Tensor] | None: ... +def get(identifier: str | dict[str, Any] | Regularizer) -> Regularizer: ... def __getattr__(name: str) -> Any: ... From b3b3e058d7fc3b7b1dd229abcd476211e48015b3 Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Wed, 22 Feb 2023 00:31:41 -0800 Subject: [PATCH 7/9] adjust order of 2 overloads for overload overlap rule. --- stubs/tensorflow/tensorflow/keras/constraints.pyi | 4 ++-- stubs/tensorflow/tensorflow/keras/regularizers.pyi | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/constraints.pyi b/stubs/tensorflow/tensorflow/keras/constraints.pyi index 4ef962cdf835..2a8c4d36eae7 100644 --- a/stubs/tensorflow/tensorflow/keras/constraints.pyi +++ b/stubs/tensorflow/tensorflow/keras/constraints.pyi @@ -14,7 +14,7 @@ _Constraint: TypeAlias = str | dict[str, Any] | Constraint | None # noqa: Y047 @overload def get(identifier: None) -> None: ... @overload -def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... -@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: ... diff --git a/stubs/tensorflow/tensorflow/keras/regularizers.pyi b/stubs/tensorflow/tensorflow/keras/regularizers.pyi index e73df6b0eda1..5d35f6c3eefb 100644 --- a/stubs/tensorflow/tensorflow/keras/regularizers.pyi +++ b/stubs/tensorflow/tensorflow/keras/regularizers.pyi @@ -15,7 +15,7 @@ _Regularizer: TypeAlias = str | dict[str, Any] | Regularizer | None # noqa: Y04 @overload def get(identifier: None) -> None: ... @overload -def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... -@overload def get(identifier: str | dict[str, Any] | Regularizer) -> Regularizer: ... +@overload +def get(identifier: Callable[[Tensor], Tensor]) -> Callable[[Tensor], Tensor]: ... def __getattr__(name: str) -> Any: ... From 068a53041eae0445e5da5e531fa8115289d69509 Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Wed, 8 Mar 2023 23:33:21 -0800 Subject: [PATCH 8/9] Treat build/compute output shape positional only. --- stubs/tensorflow/@tests/stubtest_allowlist.txt | 7 +++++++ stubs/tensorflow/tensorflow/keras/layers.pyi | 11 ++++------- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/stubs/tensorflow/@tests/stubtest_allowlist.txt b/stubs/tensorflow/@tests/stubtest_allowlist.txt index febfb497cfd2..3d6d25232256 100644 --- a/stubs/tensorflow/@tests/stubtest_allowlist.txt +++ b/stubs/tensorflow/@tests/stubtest_allowlist.txt @@ -57,3 +57,10 @@ 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 diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi index 993a3298e6e9..e076440ecbe9 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -60,11 +60,11 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): # input_shape's real type depends on _InputT, but we can't express that without HKT. # For example _InputT tf.Tensor -> tf.TensorShape, _InputT dict[str, tf.Tensor] -> dict[str, tf.TensorShape]. - def build(self, input_shape: Any) -> None: ... + def build(self, input_shape: Any, /) -> None: ... @overload - def compute_output_shape(self: Layer[tf.Tensor, tf.Tensor], input_shape: tf.TensorShape) -> tf.TensorShape: ... + def compute_output_shape(self: Layer[tf.Tensor, tf.Tensor], input_shape: tf.TensorShape, /) -> tf.TensorShape: ... @overload - def compute_output_shape(self, input_shape: Any) -> Any: ... + def compute_output_shape(self, input_shape: Any, /) -> Any: ... def add_weight( self, name: str | None = None, @@ -110,7 +110,7 @@ _LayerDtype: TypeAlias = tf._DTypeLike | dict[str, Any] | Any # Layer's compute_output_shape commonly have instance as first argument name instead of self. # This is an artifact of actual implementation commonly uses a decorator to define it. -# Layer.build has same weirdness sometimes. +# Layer.build has same weirdness sometimes. For both marked as positional only. class Dense(Layer[tf.Tensor, tf.Tensor]): def __init__( self, @@ -163,7 +163,6 @@ class ReLU(Layer[tf.Tensor, tf.Tensor]): dynamic: bool = False, name: str | None = None, ) -> None: ... - def compute_output_shape(instance, input_shape: tf.TensorShape) -> tf.TensorShape: ... class Dropout(Layer[tf.Tensor, tf.Tensor]): def __init__( @@ -192,7 +191,5 @@ class Embedding(Layer[tf.Tensor, tf.Tensor]): dynamic: bool = False, name: str | None = None, ) -> None: ... - def build(instance, input_shape: tf.TensorShape) -> None: ... - def compute_output_shape(instance, input_shape: tf.TensorShape) -> tf.TensorShape: ... def __getattr__(name: str) -> Incomplete: ... From 0d0ad27112f1d91030b3ee4e46761e75411a0e7f Mon Sep 17 00:00:00 2001 From: Mehdi Drissi Date: Thu, 9 Mar 2023 00:22:41 -0800 Subject: [PATCH 9/9] address remaining comments --- stubs/tensorflow/tensorflow/keras/constraints.pyi | 3 --- stubs/tensorflow/tensorflow/keras/layers.pyi | 10 ++++++---- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/stubs/tensorflow/tensorflow/keras/constraints.pyi b/stubs/tensorflow/tensorflow/keras/constraints.pyi index 2a8c4d36eae7..12ba257cb66b 100644 --- a/stubs/tensorflow/tensorflow/keras/constraints.pyi +++ b/stubs/tensorflow/tensorflow/keras/constraints.pyi @@ -1,7 +1,6 @@ from _typeshed import Incomplete from collections.abc import Callable from typing import Any, overload -from typing_extensions import TypeAlias from tensorflow import Tensor @@ -9,8 +8,6 @@ class Constraint: def get_config(self) -> dict[str, Any]: ... def __call__(self, __w: Tensor) -> Tensor: ... -_Constraint: TypeAlias = str | dict[str, Any] | Constraint | None # noqa: Y047 - @overload def get(identifier: None) -> None: ... @overload diff --git a/stubs/tensorflow/tensorflow/keras/layers.pyi b/stubs/tensorflow/tensorflow/keras/layers.pyi index e076440ecbe9..7bb643dd9741 100644 --- a/stubs/tensorflow/tensorflow/keras/layers.pyi +++ b/stubs/tensorflow/tensorflow/keras/layers.pyi @@ -7,7 +7,7 @@ import tensorflow as tf from tensorflow import Tensor, Variable, VariableAggregation, VariableSynchronization, _TensorCompatible from tensorflow._aliases import AnyArray from tensorflow.keras.activations import _Activation -from tensorflow.keras.constraints import _Constraint +from tensorflow.keras.constraints import Constraint from tensorflow.keras.initializers import _Initializer from tensorflow.keras.regularizers import _Regularizer @@ -60,11 +60,11 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): # input_shape's real type depends on _InputT, but we can't express that without HKT. # For example _InputT tf.Tensor -> tf.TensorShape, _InputT dict[str, tf.Tensor] -> dict[str, tf.TensorShape]. - def build(self, input_shape: Any, /) -> None: ... + def build(self, __input_shape: Any) -> None: ... @overload - def compute_output_shape(self: Layer[tf.Tensor, tf.Tensor], input_shape: tf.TensorShape, /) -> tf.TensorShape: ... + def compute_output_shape(self: Layer[tf.Tensor, tf.Tensor], __input_shape: tf.TensorShape) -> tf.TensorShape: ... @overload - def compute_output_shape(self, input_shape: Any, /) -> Any: ... + def compute_output_shape(self, __input_shape: Any) -> Any: ... def add_weight( self, name: str | None = None, @@ -108,6 +108,8 @@ class Layer(Generic[_InputT, _OutputT], tf.Module): # TODO: Replace last Any after adding tf.keras.mixed_precision.Policy. _LayerDtype: TypeAlias = tf._DTypeLike | dict[str, Any] | Any +_Constraint: TypeAlias = str | dict[str, Any] | Constraint | None + # Layer's compute_output_shape commonly have instance as first argument name instead of self. # This is an artifact of actual implementation commonly uses a decorator to define it. # Layer.build has same weirdness sometimes. For both marked as positional only.