From 236e04485dbf162c974cfd2d6f2922d42eb1bd8f Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 5 May 2026 13:59:46 +0200 Subject: [PATCH 1/2] fix mypy errors in `einops.layers.tensorflow` and add missing annotations --- einops/layers/tensorflow.py | 38 ++++++++++++++++++++----------------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/einops/layers/tensorflow.py b/einops/layers/tensorflow.py index 316c6e35..f08e2d1a 100644 --- a/einops/layers/tensorflow.py +++ b/einops/layers/tensorflow.py @@ -11,7 +11,7 @@ """ -from typing import cast +from typing import Any, TypeVar, cast import tensorflow as tf from tensorflow.keras.layers import Layer @@ -21,31 +21,35 @@ __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): + +class Rearrange(RearrangeMixin, Layer[_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, Layer[_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, Layer[_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 +60,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 +86,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 +96,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, From 24beb15f672accc7333cbd79f6321ee70240b7cc Mon Sep 17 00:00:00 2001 From: jorenham Date: Tue, 5 May 2026 14:13:51 +0200 Subject: [PATCH 2/2] workaround for `tensorflow.keras.layers.Layer` not being subscriptable at runtime --- einops/layers/tensorflow.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/einops/layers/tensorflow.py b/einops/layers/tensorflow.py index f08e2d1a..d97af166 100644 --- a/einops/layers/tensorflow.py +++ b/einops/layers/tensorflow.py @@ -11,7 +11,7 @@ """ -from typing import Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast import tensorflow as tf from tensorflow.keras.layers import Layer @@ -25,8 +25,16 @@ _InputT_contra = TypeVar("_InputT_contra", contravariant=True) _OutputT_co = TypeVar("_OutputT_co", covariant=True) +# `tensorflow.keras.layers.Layer` is only generic in the stubs, not at runtime. +if TYPE_CHECKING: -class Rearrange(RearrangeMixin, Layer[_InputT_contra, _OutputT_co]): + 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 @@ -37,7 +45,7 @@ def get_config(self) -> dict[str, Any]: return {"pattern": self.pattern, **self.axes_lengths} -class Reduce(ReduceMixin, Layer[_InputT_contra, _OutputT_co]): +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 @@ -48,7 +56,7 @@ def get_config(self) -> dict[str, Any]: return {"pattern": self.pattern, "reduction": self.reduction, **self.axes_lengths} -class EinMix(_EinmixMixin, Layer[_InputT_contra, _OutputT_co]): +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