Include mypy in run_all_tests

[zaxtax]

This continues work on #688 to include mypy in testing

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).
• This change is discussed in a Github issue/
  discussion (please add a
  link).
• The documentation and docstrings adhere to the
  documentation guidelines.
• This change includes necessary high-coverage tests.
  (No quality testing = no merge!)

[codecov-commenter]

Codecov Report

Merging #2670 (5b809bb) into main (09b6d56) will increase coverage by 0.02%.
The diff coverage is 76.31%.

@@            Coverage Diff             @@
##             main    #2670      +/-   ##
==========================================
+ Coverage   81.15%   81.17%   +0.02%     
==========================================
  Files          51       51              
  Lines        5492     5493       +1     
==========================================
+ Hits         4457     4459       +2     
+ Misses       1035     1034       -1     

Impacted Files	Coverage Δ
flax/core/nn/attention.py	13.47% <0.00%> (ø)
flax/core/nn/linear.py	42.85% <0.00%> (ø)
flax/io.py	85.26% <50.00%> (ø)
flax/linen/attention.py	95.96% <50.00%> (ø)
flax/core/meta.py	98.64% <66.66%> (ø)
flax/core/scope.py	90.13% <75.00%> (ø)
flax/training/checkpoints.py	66.14% <84.61%> (+0.26%)	⬆️
flax/core/frozen_dict.py	96.22% <100.00%> (ø)
flax/linen/dotgetter.py	72.91% <100.00%> (ø)
flax/linen/linear.py	97.51% <100.00%> (ø)
... and 5 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[levskaya]

Comments below.

Just curious - does modern mypy use mypyc compiler by default or is that a random side-effort?

[levskaya]

mypy doesn't support bound type-vars?! is this intentional Any cast?

[zaxtax]

mypy has had lots of trouble dealing with type-vars in general. I've tried to preserve their use, and just define them as Any. That way, we can reintroduce stricter checking later where possible.

[levskaya]

could we add a IntSequence (or something) to indicate what this Any is supposed to be for a human reader?

[zaxtax]

Sounds like a great idea! Just as an FYI, there are many places in the code where pytype wanted something to be Iterabale and mypy wanted something to be Sequence. I couldn't make both happy so they ended up with Any

[levskaya]

that's both sad and funny... I wonder if there's a way to harmonize these two in our code better or to agitate for change - for now just another Any alias indicating we intend a sequence of ints should be ok

[levskaya]

are we messing up internal type annotations in flax module methods to require this opt-out?

[levskaya]

when we add new open source config files we need to configure them to be sync'd by the internal sync system first - I just added your mypy.ini file to #2693 - maybe rebase on that after it goes in (hopefully soon).

[zaxtax]

That and mypy not being able to reason about flax's dataclasses were the two annoying bits of making this PR work. On the bright side, as things like mypy improve, we can reintroduce more stringent types.

…

On Wed, 7 Dec 2022, 11:49 Anselm Levskaya, ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In flax/linen/experimental/layers_with_named_axes.py <#2670 (comment)>: > @@ -30,10 +30,10 @@ # Type annotations Array = jnp.ndarray -Axes = Union[int, Iterable[int]] -DType = jnp.dtype +Axes = Union[int, Any] that's both sad and funny... I wonder if there's a way to harmonize these two in our code better or to agitate for change - for now just another Any alias indicating we intend a sequence of ints should be ok — Reply to this email directly, view it on GitHub <#2670 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCULM2ZAD5ADW3S4JY33WMBTUBANCNFSM6AAAAAASOVUNZI> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[levskaya]

@zaxtax - seeing a small number of new internal pytype bugs of this sort:

File "rax/examples/flax_integration/main.py", line 143, in train_step: Function FrozenDict.copy was called with the wrong arguments [wrong-arg-types]
         Expected: (self, add_or_replace: flax.core.frozen_dict.FrozenDict[str, Mapping[str, Any]])
  Actually passed: (self, add_or_replace: Dict[str, Any])

File "jestimator/states.py", line 73, in variables: Function FrozenDict.copy was called with the wrong arguments [wrong-arg-types]
         Expected: (self, add_or_replace: flax.core.frozen_dict.FrozenDict[str, Any])
  Actually passed: (self, add_or_replace: Dict[str, Optional[flax.core.frozen_dict.FrozenDict[str, Any]]])

File "jestimator/states.py", line 143, in variables: Function FrozenDict.copy was called with the wrong arguments [wrong-arg-types]
         Expected: (self, add_or_replace: flax.core.frozen_dict.FrozenDict[str, Any])
  Actually passed: (self, add_or_replace: Dict[str, Optional[flax.core.frozen_dict.FrozenDict[str, Any]]])

[zaxtax]

@zaxtax - seeing a small number of new internal pytype bugs of this sort:

File "rax/examples/flax_integration/main.py", line 143, in train_step: Function FrozenDict.copy was called with the wrong arguments [wrong-arg-types]
         Expected: (self, add_or_replace: flax.core.frozen_dict.FrozenDict[str, Mapping[str, Any]])
  Actually passed: (self, add_or_replace: Dict[str, Any])

File "jestimator/states.py", line 73, in variables: Function FrozenDict.copy was called with the wrong arguments [wrong-arg-types]
         Expected: (self, add_or_replace: flax.core.frozen_dict.FrozenDict[str, Any])
  Actually passed: (self, add_or_replace: Dict[str, Optional[flax.core.frozen_dict.FrozenDict[str, Any]]])

File "jestimator/states.py", line 143, in variables: Function FrozenDict.copy was called with the wrong arguments [wrong-arg-types]
         Expected: (self, add_or_replace: flax.core.frozen_dict.FrozenDict[str, Any])
  Actually passed: (self, add_or_replace: Dict[str, Optional[flax.core.frozen_dict.FrozenDict[str, Any]]])

@levskaya try now