google / google/flax

Better type Module.apply signature - PR discussion

Open
#4,779 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Following up on #2086 and this recent SO question [How to type hint `flax.linen.Module.apply`'s output correctly?](https://stackoverflow.com/q/79658104/12439683).

I would like to add a PR based on my [answer](https://stackoverflow.com/a/79658693/12439683) and add two `typing.overload` signatures that make it more clear when to expect a `tuple` return value - following the contribution guidelines I am opening this issue first.

Based on #2086, the docstring and

https://github.com/google/flax/blob/893a66024aad7c9ab0d0c5b6cb22578b8b95ae4c/flax/core/scope.py#L1080-L1083

I am certain that only `mutable=False` never returns a tuple, otherwise it will always be a tuple.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.