Repository: larq/larq

Issue:
Test against TensorFlow 2.0.0-rc0



Discussion:
This failure looks like it is related to to our training metrics. Though I am not sure if this is a problem with our code or with TF 2rc0.
I have a workaround and posted an issue on TensorFlow to make sure this behaviour change is intended and not a bug: https://github.com/tensorflow/tensorflow/issues/32210

Resolving patch:
diff --git a/larq/layers_base.py b/larq/layers_base.py
--- a/larq/layers_base.py
+++ b/larq/layers_base.py
@@ -48,8 +48,11 @@ def build(self, input_shape):
     def non_trainable_weights(self):
         weights = super().non_trainable_weights
         if hasattr(self, "flip_ratio"):
-            metrics_weights = self.flip_ratio.weights
-            return [weight for weight in weights if weight not in metrics_weights]
+            return [
+                weight
+                for weight in weights
+                if not any(weight is metric_w for metric_w in self.flip_ratio.weights)
+            ]
         return weights
 
     def call(self, inputs):
@@ -115,8 +118,11 @@ def build(self, input_shape):
     def non_trainable_weights(self):
         weights = super().non_trainable_weights
         if hasattr(self, "flip_ratio"):
-            metrics_weights = self.flip_ratio.weights
-            return [weight for weight in weights if weight not in metrics_weights]
+            return [
+                weight
+                for weight in weights
+                if not any(weight is metric_w for metric_w in self.flip_ratio.weights)
+            ]
         return weights
 
     def call(self, inputs):
@@ -206,7 +212,11 @@ def non_trainable_weights(self):
         if hasattr(self, "pointwise_flip_ratio"):
             metrics_weights.extend(self.pointwise_flip_ratio.weights)
         if metrics_weights:
-            return [weight for weight in weights if weight not in metrics_weights]
+            return [
+                weight
+                for weight in weights
+                if not any(weight is metric_w for metric_w in metrics_weights)
+            ]
         return weights
 
     def call(self, inputs):
diff --git a/larq/metrics.py b/larq/metrics.py
--- a/larq/metrics.py
+++ b/larq/metrics.py
@@ -118,7 +118,7 @@ def result(self):
 
     def reset_states(self):
         tf.keras.backend.batch_set_value(
-            [(v, 0) for v in self.variables if v != self._previous_values]
+            [(v, 0) for v in self.variables if v is not self._previous_values]
         )
 
     def get_config(self):
diff --git a/larq/models.py b/larq/models.py
--- a/larq/models.py
+++ b/larq/models.py
@@ -131,7 +131,7 @@ def __init__(self, layer):
             WeightProfile(
                 weight,
                 self._get_bitwidth(weight),
-                trainable=weight in layer.trainable_weights,
+                trainable=any(weight is w for w in layer.trainable_weights),
             )
             for weight in layer.weights
         ]
@@ -228,18 +228,16 @@ def generate_table_row(self, table_config):
 
         return row
 
-    def _quantized_weights(self):
-        try:
-            return self._layer.quantized_latent_weights
-        except:
-            return []
-
     def _get_bitwidth(self, weight):
         try:
-            quantizer = self._layer.quantizers[self._quantized_weights().index(weight)]
-            return quantizer.precision
+            for quantizer, quantized_weight in zip(
+                self._layer.quantizers, self._layer.quantized_latent_weights
+            ):
+                if quantized_weight is weight:
+                    return quantizer.precision
         except:
-            return 32
+            pass
+        return 32
 
 
 class ModelProfile(LayerProfile):


Test patch: