Add stubs

Author: zaxtax

flax/linen/linear.py

@@ -293,7 +293,7 @@ class _Conv(Module):
   bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = zeros
 
   @property
-  def shared_weights(self) -> bool:
+  def shared_weights(self) -> bool:  # type: ignore
     """Defines whether weights are shared or not between different pixels.
 
     Returns:
@@ -302,7 +302,7 @@ def shared_weights(self) -> bool:
       "locally connected layer", "unshared convolution", or "local convolution".
 
     """
-    ... # type: ignore
+    ...
 
   @compact
   def __call__(self, inputs: Array) -> Array:

mypy.ini

@@ -29,3 +29,5 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-jaxlib.mlir.*]
 ignore_missing_imports = True
+[mypy-yaml]
+ignore_missing_imports = True