Unpin pytype dependency and fix various pytype/typing errors.

PiperOrigin-RevId: 438008523

Author: levskaya

examples/nlp_seq/input_pipeline.py

@@ -18,7 +18,7 @@
 import collections
 import enum
 
-import tensorflow.compat.v2 as tf
+import tensorflow.compat.v2 as tf  # pytype: disable=import-error
 
 
 # Values for padding, unknown words and a root.

flax/linen/linear.py

@@ -205,7 +205,7 @@ def _conv_dimension_numbers(input_shape):
   return lax.ConvDimensionNumbers(lhs_spec, rhs_spec, out_spec)
 
 
-PaddingLike = Union[str, int, Sequence[Union[int, Tuple[int, int]]]] 
+PaddingLike = Union[str, int, Sequence[Union[int, Tuple[int, int]]]]
 LaxPadding = Union[str, Sequence[Tuple[int, int]]]
 
 
@@ -281,7 +281,6 @@ class _Conv(Module):
   bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = zeros
 
   @property
-  @abc.abstractmethod
   def shared_weights(self) -> bool:
     """Defines whether weights are shared or not between different pixels.
 

flax/linen/module.py

@@ -22,7 +22,7 @@
 import typing
 from typing import (Any, Callable, Dict, Iterable, List, Optional,
                     Set, Tuple, Type, TypeVar, Union, overload)
-from typing_extensions import dataclass_transform # pytype: disable=import-error
+from typing_extensions import dataclass_transform  # pytype: disable=not-supported-yet
 import weakref
 
 from flax import config

flax/linen/recurrent.py

@@ -46,7 +46,6 @@ class RNNCellBase(Module):
   """RNN cell base class."""
 
   @staticmethod
-  @abc.abstractmethod
   def initialize_carry(rng, batch_dims, size, init_fn=zeros):
     """Initialize the RNN cell carry.
 

flax/metrics/tensorboard.py

@@ -20,7 +20,7 @@
 # pylint: disable=g-import-not-at-top
 import numpy as np
 
-import tensorflow.compat.v2 as tf
+import tensorflow.compat.v2 as tf  # pytype: disable=import-error
 from tensorboard.plugins.hparams import api as hparams_api
 
 

flax/optim/adadelta.py

@@ -14,6 +14,7 @@
 
 """Adadelta Optimizer."""
 
+from typing import Optional
 from .. import struct
 from .base import OptimizerDef
 import jax.numpy as jnp
@@ -43,7 +44,7 @@ class Adadelta(OptimizerDef):
   """
 
   def __init__(self,
-               learning_rate: float = None,
+               learning_rate: Optional[float] = None,
                rho: float = 0.9,
                eps: float = 1e-6,
                weight_decay: float = 0.0):

flax/optim/adagrad.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 import jax.numpy as jnp
 import numpy as np
 from .. import struct
@@ -35,9 +36,9 @@ class _AdagradParamState:
 
 class Adagrad(OptimizerDef):
   """Adagrad optimizer"""
-  def __init__(self, learning_rate: float = None, eps=1e-8):
+  def __init__(self, learning_rate: Optional[float] = None, eps=1e-8):
     """Constructor for the Adagrad optimizer.
-        
+
     Args:
       learning_rate: the step size used to update the parameters.
     """

flax/optim/base.py

@@ -15,7 +15,7 @@
 """Flax Optimizer api."""
 
 import dataclasses
-from typing import Any, List, Tuple
+from typing import Any, List, Tuple, Optional
 import warnings
 
 from .. import jax_utils
@@ -40,7 +40,7 @@ class OptimizerState:
 
 class OptimizerDef:
   """Base class for an optimizer defintion, which specifies the initialization and gradient application logic.
-  
+
   See docstring of :class:`Optimizer` for more details.
   """
 
@@ -122,7 +122,7 @@ def update_hyper_params(self, **hyper_param_overrides):
       hp = hp.replace(**hyper_param_overrides)
     return hp
 
-  def create(self, target, focus: 'ModelParamTraversal' = None):
+  def create(self, target, focus: Optional['ModelParamTraversal'] = None):
     """Creates a new optimizer for the given target.
 
     See docstring of :class:`Optimizer` for more details.
@@ -133,7 +133,7 @@ def create(self, target, focus: 'ModelParamTraversal' = None):
         of variables dicts, e.g. `(v1, v2)` and  `('var1': v1, 'var2': v2)`
         are valid inputs as well.
       focus: a `flax.traverse_util.Traversal` that selects which subset of
-        the target is optimized. See docstring of :class:`MultiOptimizer` 
+        the target is optimized. See docstring of :class:`MultiOptimizer`
         for an example of how to define a `Traversal` object.
     Returns:
       An instance of `Optimizer`.
@@ -183,10 +183,10 @@ class _NoAux:
 class Optimizer(struct.PyTreeNode):
   """
   Flax optimizers are created using the :class:`OptimizerDef` class. That class
-  specifies the initialization and gradient application logic. Creating an 
-  optimizer using the :meth:`OptimizerDef.create` method will result in an 
+  specifies the initialization and gradient application logic. Creating an
+  optimizer using the :meth:`OptimizerDef.create` method will result in an
   instance of the :class:`Optimizer` class, which encapsulates the optimization
-  target and state. The optimizer is updated using the method 
+  target and state. The optimizer is updated using the method
   :meth:`apply_gradient`.
 
   Example of constructing an optimizer for a model::

flax/optim/rmsprop.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 import jax.numpy as jnp
 import numpy as np
 from .. import struct
@@ -38,7 +39,7 @@ class _RMSPropParamState:
 
 class RMSProp(OptimizerDef):
   """RMSProp optimizer"""
-  def __init__(self, learning_rate: float = None, beta2=0.9, eps=1e-8,
+  def __init__(self, learning_rate: Optional[float] = None, beta2=0.9, eps=1e-8,
                centered=False):
     """Constructor for the RMSProp optimizer
 
@@ -73,7 +74,7 @@ def apply_param_gradient(self, step, hyper_params, param, state, grad):
     else:
       new_mg = state.mg
       maybe_centered_v = new_v
-    new_param = param - hyper_params.learning_rate * grad / ( 
+    new_param = param - hyper_params.learning_rate * grad / (
         jnp.sqrt(maybe_centered_v) + hyper_params.eps)
     new_state = _RMSPropParamState(new_v, new_mg)
 

flax/struct.py

@@ -22,7 +22,7 @@
 from . import serialization
 
 import jax
-from typing_extensions import dataclass_transform  # pytype: disable=import-error
+from typing_extensions import dataclass_transform  # pytype: disable=not-supported-yet
 
 
 _T = TypeVar("_T")

flax/training/checkpoints.py

@@ -28,7 +28,7 @@
 from flax import core
 from flax import errors
 from flax import serialization
-from tensorflow.io import gfile
+from tensorflow.io import gfile  # pytype: disable=import-error
 
 
 # Single-group reg-exps for int or float numerical substrings.

pytype.cfg

@@ -0,0 +1,8 @@
+# NOTE: All relative paths are relative to github root directory.
+
+[pytype]
+
+# TODO(levskaya): figure out why we get pyi-error from flax's root __init__.py
+# could be a pytype bug.
+disable =
+    pyi-error

setup.py

@@ -44,7 +44,7 @@
     "pytest",
     "pytest-cov",
     "pytest-xdist==1.34.0",  # upgrading to 2.0 broke tests, need to investigate
-    "pytype==2021.5.25",  # pytype 2021.6.17 complains on recurrent.py, need to investigate!
+    "pytype",
     "sentencepiece",  # WMT example.
     "svn",
     "tensorflow_text>=2.4.0",  # WMT example.

tests/run_all_tests.sh

@@ -51,35 +51,19 @@ pytest -n auto tests $PYTEST_OPTS
 # In pytest foo/bar/baz_test.py and baz/bleep/baz_test.py will collide and error out when
 # /foo/bar and /baz/bleep aren't set up as packages.
 for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
-    # Skip tests on deprecated example -- we get this error.
-    #   ValueError: Failed to construct dataset lm1b: BuilderConfig subwords32k not found. Available: []
-    #
-    # TODO: Remove after github.com/google/flax/issues/567 is resolved
-    if [[ "$egd" == "examples/lm1b_deprecated" ]]; then
-        continue
-    fi
     pytest $egd
 done
 
-# validate types
-if [[ "$OSTYPE" == "darwin"* ]]; then
-    echo "Pytype is currently not working on MacOS, see https://github.com/google/pytype/issues/661"
-else
-    pytype flax/
-
-    for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
-        # Skip pytype on deprecated example
-        # TODO: Remove after github.com/google/flax/issues/567 is resolved
-        if [[ "$egd" == "examples/lm1b_deprecated" ]]; then
-            continue
-        fi
+# Validate types in library code.
+pytype --config pytype.cfg flax/
 
-        # use cd to make sure pytpe cache lives in example dir and doesn't name clash
-        # use *.py to avoid importing configs as a top-level import which leads tot import errors
-        # because config files use relative imports (e.g. from config import ...).
-        (cd $egd ; pytype "*.py")
-    done
-fi
+# Validate types in examples.
+for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
+    # use cd to make sure pytpe cache lives in example dir and doesn't name clash
+    # use *.py to avoid importing configs as a top-level import which leads to import errors
+    # because config files use relative imports (e.g. from config import ...).
+    (cd $egd ; pytype --config ../../pytype.cfg "*.py")
+done
 
 # Return error code 0 if no real failures happened.
 echo "finished all tests."