fix mypy errors in einops.layers.tensorflow and add missing annotations

einops/layers/tensorflow.py

@@ -11,7 +11,7 @@
 
 """
 
-from typing import cast
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast
 
 import tensorflow as tf
 from tensorflow.keras.layers import Layer
@@ -21,31 +21,43 @@
 
 __author__ = "Alex Rogozhnikov"
 
+# matches the type parameters used by `tensorflow.keras.layers.Layer`
+_InputT_contra = TypeVar("_InputT_contra", contravariant=True)
+_OutputT_co = TypeVar("_OutputT_co", covariant=True)
 
-class Rearrange(RearrangeMixin, Layer):
-    def build(self, input_shape):
+# `tensorflow.keras.layers.Layer` is only generic in the stubs, not at runtime.
+if TYPE_CHECKING:
+
+    class _BaseLayer(Layer[_InputT_contra, _OutputT_co], Generic[_InputT_contra, _OutputT_co]): ...
+else:
+
+    class _BaseLayer(Layer, Generic[_InputT_contra, _OutputT_co]): ...
+
+
+class Rearrange(RearrangeMixin, _BaseLayer[_InputT_contra, _OutputT_co]):
+    def build(self, input_shape: Any) -> None:
         pass  # layer does not have any parameters to be initialized
 
-    def call(self, inputs):
+    def call(self, inputs: _InputT_contra) -> _OutputT_co:
         return self._apply_recipe(inputs)
 
-    def get_config(self):
+    def get_config(self) -> dict[str, Any]:
         return {"pattern": self.pattern, **self.axes_lengths}
 
 
-class Reduce(ReduceMixin, Layer):
-    def build(self, input_shape):
+class Reduce(ReduceMixin, _BaseLayer[_InputT_contra, _OutputT_co]):
+    def build(self, input_shape: Any) -> None:
         pass  # layer does not have any parameters to be initialized
 
-    def call(self, inputs):
+    def call(self, inputs: _InputT_contra) -> _OutputT_co:
         return self._apply_recipe(inputs)
 
-    def get_config(self):
+    def get_config(self) -> dict[str, Any]:
         return {"pattern": self.pattern, "reduction": self.reduction, **self.axes_lengths}
 
 
-class EinMix(_EinmixMixin, Layer):
-    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound):
+class EinMix(_EinmixMixin, _BaseLayer[_InputT_contra, _OutputT_co]):
+    def _create_parameters(self, weight_shape, weight_bound, bias_shape, bias_bound) -> None:
         # this method is called in __init__,
         #  but we postpone actual creation to build(), as TF instruction suggests
         self._params = [weight_shape, weight_bound, bias_shape, bias_bound]
@@ -56,16 +68,16 @@ def _create_rearrange_layers(
         pre_reshape_lengths: dict | None,
         post_reshape_pattern: str | None,
         post_reshape_lengths: dict | None,
-    ):
-        self.pre_rearrange = None
+    ) -> None:
+        self.pre_rearrange: Rearrange | None = None
         if pre_reshape_pattern is not None:
             self.pre_rearrange = Rearrange(pre_reshape_pattern, **cast(dict, pre_reshape_lengths))
 
-        self.post_rearrange = None
+        self.post_rearrange: Rearrange | None = None
         if post_reshape_pattern is not None:
             self.post_rearrange = Rearrange(post_reshape_pattern, **cast(dict, post_reshape_lengths))
 
-    def build(self, input_shape):
+    def build(self, input_shape: Any) -> None:
         [weight_shape, weight_bound, bias_shape, bias_bound] = self._params
         self.weight = self.add_weight(
             shape=weight_shape,
@@ -82,7 +94,7 @@ def build(self, input_shape):
         else:
             self.bias = None
 
-    def call(self, inputs):
+    def call(self, inputs: _InputT_contra) -> _OutputT_co:
         if self.pre_rearrange is not None:
             inputs = self.pre_rearrange(inputs)
         result = tf.einsum(self.einsum_pattern, inputs, self.weight)
@@ -92,7 +104,7 @@ def call(self, inputs):
             result = self.post_rearrange(result)
         return result
 
-    def get_config(self):
+    def get_config(self) -> dict[str, Any]:
         return {
             "pattern": self.pattern,
             "weight_shape": self.weight_shape,