#303 fix dtype policy bug in GEM layers. (#304)

* GEM layers create a general pooling layer in the init, but we didn't
pass the kwargs. This means the general pooling layer didn't have the
dtype policy. This caused the GEM layers to fail when using a
mixed_float dtype policy as the general pooling layer returns float32
and the GEM dtype policy is float16.

The fix is to pass all kwargs onto the general pooling layer.

* Patch bump

* Cap the TF version at 2.9 for the current master branch.

Author: Owen Vallis

setup.py

@@ -78,7 +78,7 @@ def get_version(rel_path):
             "mkdocs-autorefs",
             "mkdocs-material",
             "mkdocstrings",
-            "mypy",
+            "mypy<=0.982",
             "pytest",
             "pytype",
             "setuptools",
@@ -87,9 +87,9 @@ def get_version(rel_path):
             "types-tabulate",
             "wheel",
         ],
-        "tensorflow": ["tensorflow>=2.4"],
-        "tensorflow-gpu": ["tensorflow-gpu>=2.4"],
-        "tensorflow-cpu": ["tensorflow-cpu>=2.4"],
+        "tensorflow": ["tensorflow>=2.4,<=2.9"],
+        "tensorflow-gpu": ["tensorflow-gpu>=2.4,<=2.9"],
+        "tensorflow-cpu": ["tensorflow-cpu>=2.4,<=2.9"],
     },
     classifiers=[
         "Development Status :: 5 - Production/Stable",

tensorflow_similarity/__init__.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-__version__ = "0.16.8"
+__version__ = "0.16.9"
 
 
 from . import algebra  # noqa

tensorflow_similarity/layers.py

@@ -145,7 +145,7 @@ def __init__(self, p: float = 3.0, data_format: Optional[str] = None, keepdims:
         super().__init__(p=p, data_format=data_format, keepdims=keepdims, **kwargs)
 
         self.input_spec = layers.InputSpec(ndim=3)
-        self.gap = layers.GlobalAveragePooling1D(data_format=data_format, keepdims=keepdims)
+        self.gap = layers.GlobalAveragePooling1D(data_format=data_format, keepdims=keepdims, **kwargs)
         self.step_axis = 1 if self.data_format == "channels_last" else 2
 
     def call(self, inputs: FloatTensor) -> FloatTensor:
@@ -231,7 +231,7 @@ def __init__(self, p: float = 3.0, data_format: Optional[str] = None, keepdims:
         super().__init__(p=p, data_format=data_format, keepdims=keepdims, **kwargs)
 
         self.input_spec = layers.InputSpec(ndim=4)
-        self.gap = layers.GlobalAveragePooling2D(data_format, keepdims)
+        self.gap = layers.GlobalAveragePooling2D(data_format=data_format, keepdims=keepdims, **kwargs)
 
     def call(self, inputs: FloatTensor) -> FloatTensor:
         x = inputs