[hmc-cs-mdrissi]

Summary

Next tensorflow stub pr. Primary goal of this one was to define a reasonable initial stub for Layers, one of the most common types user interacts with. The other files are mainly types needed to cover Layer's arguments. I've included a couple layer's too. Adding a stub for 1 layer is fairly easy as most of the time only thing needed is constructor and all public methods come from parent Layer class.

I ended up adding a fair amount of stubtest allowlist and tried to give an explanation for each. call/__call__ strangeness is probably biggest one and I can explain more if it's unclear. At it's core the magic is as a user, you define

def call(self, inputs): ...

but you are never supposed to call that method and instead use __call__ which does a good amount of signature introspection/argument adjustment for you.

[hmc-cs-mdrissi]

I don't think we have PEP 692 tracker ticket yet. Will make one later if no one covers it first.

[JelleZijlstra]

I think PEP 692 might already have sufficient support to use it; at least mypy and pyright support it. Haven't checked pytype though.

[hmc-cs-mdrissi]

Testing basic 692 support:

class Foo(TypedDict, total=False):
    a: int
    b: str


def f(**kwargs: Unpack[Foo]) -> None:
    ...


f(a=1, b="2")  # OK
f(a=1, b=2)  # Error: b has type str
f(a=1, b="2", c=3)  # Error: unexpected keyword argument "c"

• pyright works fine
• mypy works if we add --enable-incomplete-feature=Unpack (otherwise it gives error: "Unpack" support is experimental, use --enable-incomplete-feature=Unpack to enable)
• pytype produces typing_extensions.Unpack not supported yet [not-supported-yet]. No errors otherwise.
• pyre is same as pytype and just produces pep_692.py:11:16 Undefined or invalid type [11]: Annotation Unpack is not defined as a type.

I'll open ticket to track.

[hmc-cs-mdrissi]

This is probably most cursed weirdness of imports (pylint is unhappy too with this). A lot of other tf modules do similar thing and this part of allowlist will grow over time.

[hmc-cs-mdrissi]

Hmm optuna hit is valid, but input_dim argument is undocumented at least on layer documentation. input_dim is one of the internal arguments here.

Same issue for streamlit error. input_shape argument is also not documented anywhere I can find. There's 6 more internal layer arguments that could also be added.

I weakly prefer not to include arguments that do not appear in documentation, but I'm fine adding them in if desired.

[hmc-cs-mdrissi]

General review poke. There are a few small questions I’ve left but otherwise been ready for over week.

[JelleZijlstra]

Using the type alias seems unnecessary here as all types but Constraint are already covered by previous overloads.

[hmc-cs-mdrissi]

Fixed and made overloads non overlapping.

One question for this and similar comments is it recommended to make these generic? As,

from tensorflow.initializers import Zeros

get(Zeros()) # returns same type Zeros().

I feel neutral as most common overload I think is str/dict[str, Any] case but otherwise constraint case always returns same type.

[JelleZijlstra]

Same as before, seems like this is only the Callable.

[hmc-cs-mdrissi]

Fixed. For this one re-reading source I missed one case (type[Initializer] is also allowed).

[JelleZijlstra]

Sorry for the wait! I left a few comments.

[hmc-cs-mdrissi]

I think I've addressed all comments now.

[github-actions]

Diff from mypy_primer, showing the effect of this PR on open source code:

optuna (https://github.com/optuna/optuna)
+ tests/integration_tests/test_tfkeras.py:20: error: Unexpected keyword argument "input_dim" for "Dense"  [call-arg]

streamlit (https://github.com/streamlit/streamlit)
+ lib/tests/streamlit/runtime/legacy_caching/hashing_test.py: note: In member "test_tf_saved_model" of class "HashTest":
+ lib/tests/streamlit/runtime/legacy_caching/hashing_test.py:417:17: error: Unexpected keyword argument "input_shape" for "Dense"  [call-arg]
- lib/tests/streamlit/runtime/legacy_caching/hashing_test.py: note: At top level: