fix mypy + add .git-blame-ignore-revs

google · Jul 21, 2023 · 97d038c · 97d038c
1 parent 40a6e07
commit 97d038c
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 22 deletions.
diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs
@@ -0,0 +1,2 @@
+# apply pyink
+40a6e074e5224d733f964be00e21e0a1cb98bd2e
diff --git a/flax/linen/attention.py b/flax/linen/attention.py
@@ -381,10 +381,10 @@ class SelfAttention(MultiHeadDotProductAttention):
   """Self-attention special case of multi-head dot-product attention."""
 
   @compact
-  def __call__(
+  def __call__(  # type: ignore
       self,
       inputs_q: Array,
-      mask: Optional[Array] = None,  # type: ignore
+      mask: Optional[Array] = None,
       deterministic: Optional[bool] = None,
   ):
     """Applies multi-head dot product self-attention on the input data.

diff --git a/flax/linen/module.py b/flax/linen/module.py
@@ -828,13 +828,11 @@ def _customized_dataclass_transform(cls, kw_only: bool):
         for name, annotation, default in extra_fields:  # pytype: disable=invalid-annotation
           setattr(cls, name, default)
           cls.__annotations__[name] = annotation
-        dataclasses.dataclass(
+        dataclasses.dataclass(  # type: ignore[call-overload]
             unsafe_hash='__hash__' not in cls.__dict__,
             repr=False,
             kw_only=True,
-        )(
-            cls
-        )  # type: ignore[call-overload]
+        )(cls)
       else:
         raise TypeError('`kw_only` is not available before Py 3.10.')
     else:
@@ -1900,8 +1898,8 @@ def sow(
       name: str,
       value: T,
       reduce_fn: Callable[[K, T], K] = tuple_reduce,
-      init_fn: Callable[[], K] = tuple_init,
-  ) -> bool:  # type: ignore
+      init_fn: Callable[[], K] = tuple_init,  # type: ignore
+  ) -> bool:
     ...
 
   def sow(
@@ -1910,8 +1908,8 @@ def sow(
       name: str,
       value: T,
       reduce_fn: Callable[[K, T], K] = tuple_reduce,
-      init_fn: Callable[[], K] = tuple_init,
-  ) -> bool:  # type: ignore
+      init_fn: Callable[[], K] = tuple_init,  # type: ignore
+  ) -> bool:
     """Stores a value in a collection.
 
     Collections can be used to collect intermediate values without

diff --git a/flax/linen/recurrent.py b/flax/linen/recurrent.py
@@ -342,20 +342,16 @@ def _concat_dense(
           param_dtype=self.param_dtype,
           kernel_init=self.kernel_init,
           bias_init=self.bias_init,
-          name=f'i{component}',
-      )(
-          inputs
-      )  # type: ignore[call-arg]
+          name=f'i{component}',  # type: ignore[call-arg]
+      )(inputs)
       dense_params_h[component] = DenseParams(
           features=hidden_features,
           use_bias=True,
           param_dtype=self.param_dtype,
           kernel_init=self.recurrent_kernel_init,
           bias_init=self.bias_init,
-          name=f'h{component}',
-      )(
-          h
-      )  # type: ignore[call-arg]
+          name=f'h{component}',  # type: ignore[call-arg]
+      )(h)
     dense_h = _concat_dense(h, dense_params_h, use_bias=True)
     dense_i = _concat_dense(inputs, dense_params_i, use_bias=False)
 
@@ -809,8 +805,8 @@ def __call__(
     if reverse:
       inputs = jax.tree_map(
           lambda x: flip_sequences(
-              x, seq_lengths, num_batch_dims=len(batch_dims), time_major=time_major
-          ),  # type: ignore
+              x, seq_lengths, num_batch_dims=len(batch_dims), time_major=time_major  # type: ignore
+          ),
           inputs,
       )
 
@@ -867,8 +863,8 @@ def scan_fn(
     if reverse and keep_order:
       outputs = jax.tree_map(
           lambda x: flip_sequences(
-              x, seq_lengths, num_batch_dims=len(batch_dims), time_major=time_major
-          ),  # type: ignore
+              x, seq_lengths, num_batch_dims=len(batch_dims), time_major=time_major  # type: ignore
+          ),
           outputs,
       )