Repository: larq/larq

Issue:
Use metaclass to allow direct instantiation of QuantizedVariable
This allows us to remove the `create_quantized_variable` function at the expense of a slightly more complicated implementation. Users can now directly instantiate `QantizedVariable` and don't have to worry about correctly handling distribution strategies.


Discussion:

Resolving patch:
diff --git a/larq/layers_base.py b/larq/layers_base.py
--- a/larq/layers_base.py
+++ b/larq/layers_base.py
@@ -2,8 +2,8 @@
 
 import tensorflow as tf
 
-from larq import metrics as lq_metrics
-from larq import quantized_scope, quantized_variable, quantizers
+from larq import metrics as lq_metrics, quantized_scope, quantizers
+from larq.quantized_variable import QuantizedVariable
 
 log = logging.getLogger(__name__)
 
@@ -27,7 +27,7 @@ def _add_variable_with_custom_getter(self, name, **kwargs):
         # Wrap `getter` with a version that returns a `QuantizedVariable`.
         def getter(*args, **kwargs):
             variable = old_getter(*args, **kwargs)
-            return quantized_variable.create_quantized_variable(variable, quantizer)
+            return QuantizedVariable.from_variable(variable, quantizer)
 
         return super()._add_variable_with_custom_getter(name, getter=getter, **kwargs)
 
diff --git a/larq/quantized_variable.py b/larq/quantized_variable.py
--- a/larq/quantized_variable.py
+++ b/larq/quantized_variable.py
@@ -39,6 +39,63 @@ def __init__(self, variable, quantizer=None, precision=None):
         self.quantizer = quantizer
         self.precision = precision or getattr(quantizer, "precision", None)
 
+    @classmethod
+    def from_variable(cls, variable, quantizer=None, precision=None):
+        """Creates a QuantizedVariable that wraps another variable.
+
+        This typically just returns `QuantizedVariable(variable)`. But, if the variable
+        is a DistributedVariable or one of its subclasses, we instead dynamically
+        create a class that subclasses from both QuantizedVariable and
+        variable.__class__. This is so the returned variable will still pass
+        `isinstance(variable, variable.__class__)`, which is required for
+        DistributedVariables and its subclasses to work properly.
+
+        # Arguments
+        variable: A floating-point resource variable to wrap.
+        quantizer: An optional quantizer to transform the floating-point variable to a
+            fake quantized variable.
+        precision: An optional integer defining the precision of the quantized variable.
+            If `None`, `quantizer.precision` is used.
+
+        # Returns
+        A QuantizedVariable that wraps the variable.
+        """
+        if not isinstance(variable, DistributedVariable):  # type: ignore
+            return cls(variable, quantizer, precision)
+
+        class QuantizedDistributedVariable(cls, variable.__class__):
+            """A QuantizedVariable that also subclasses from DistributedVariable."""
+
+            def get(self, *args, **kwargs):
+                # For some reason this is needed to make unit `x + x` pass on TF 1.14
+                return self._quantize(self.latent_variable.get(*args, **kwargs))
+
+        return QuantizedDistributedVariable(variable, quantizer, precision)
+
+    @staticmethod
+    def _maybe_wrap(variable, quantizer, precision, wrap=True):
+        """Creates an QuantizedVariable that wraps another variable if applicable.
+
+        This function is used to wrap the return value of QuantizedVariable.assign.
+        Unfortunately MirroredVariable.assign will (incorrectly) return a Mirrored
+        value instead of a MirroredVariable. So we cannot properly wrap it in an
+        AutoCastVariable. We return the original variable in that case.
+
+        # Arguments
+        variable: A tf.Variable or op.
+        quantizer: An optional quantizer to transform the floating-point variable to a
+            fake quantized variable.
+        precision: An optional integer defining the precision of the quantized variable.
+            If `None`, `quantizer.precision` is used.
+        wrap: A boolean to define whether to wrap the variable in an QuantizedVariable.
+
+        # Returns
+        An QuantizedVariable if wrap is True and variable is a resource variable.
+        """
+        if wrap and resource_variable_ops.is_resource_variable(variable):
+            return QuantizedVariable.from_variable(variable, quantizer, precision)
+        return variable
+
     def _quantize(self, value):
         if self.quantizer and quantized_scope.should_quantize():
             return self.quantizer(value)
@@ -134,59 +191,59 @@ def constraint(self):
 
     def assign(self, value, use_locking=None, name=None, read_value=True):
         op = self.latent_variable.assign(value, use_locking, name, read_value)
-        return _maybe_wrap(op, self.quantizer, self.precision, wrap=read_value)
+        return self._maybe_wrap(op, self.quantizer, self.precision, wrap=read_value)
 
     def assign_add(self, delta, use_locking=None, name=None, read_value=True):
         op = self.latent_variable.assign_add(delta, use_locking, name, read_value)
-        return _maybe_wrap(op, self.quantizer, self.precision, wrap=read_value)
+        return self._maybe_wrap(op, self.quantizer, self.precision, wrap=read_value)
 
     def assign_sub(self, delta, use_locking=None, name=None, read_value=True):
         op = self.latent_variable.assign_sub(delta, use_locking, name, read_value)
-        return _maybe_wrap(op, self.quantizer, self.precision, wrap=read_value)
+        return self._maybe_wrap(op, self.quantizer, self.precision, wrap=read_value)
 
     def scatter_sub(self, *args, **kwargs):
         var = self.latent_variable.scatter_sub(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_add(self, *args, **kwargs):
         var = self.latent_variable.scatter_add(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_max(self, *args, **kwargs):
         var = self.latent_variable.scatter_max(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_min(self, *args, **kwargs):
         var = self.latent_variable.scatter_min(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_mul(self, *args, **kwargs):
         var = self.latent_variable.scatter_mul(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_div(self, *args, **kwargs):
         var = self.latent_variable.scatter_div(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_update(self, *args, **kwargs):
         var = self.latent_variable.scatter_update(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def batch_scatter_update(self, *args, **kwargs):
         var = self.latent_variable.batch_scatter_update(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_nd_sub(self, *args, **kwargs):
         var = self.latent_variable.scatter_nd_sub(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_nd_add(self, *args, **kwargs):
         var = self.latent_variable.scatter_nd_add(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def scatter_nd_update(self, *args, **kwargs):
         var = self.latent_variable.scatter_nd_update(*args, **kwargs)
-        return _maybe_wrap(var, self.quantizer, self.precision)
+        return self._maybe_wrap(var, self.quantizer, self.precision)
 
     def count_up_to(self, *args, **kwargs):
         return self.latent_variable.count_up_to(*args, **kwargs)
@@ -252,60 +309,3 @@ def _as_graph_element(self):
     QuantizedVariable, QuantizedVariable._dense_var_to_tensor
 )
 ops.register_dense_tensor_like_type(QuantizedVariable)
-
-
-def create_quantized_variable(variable, quantizer=None, precision=None):
-    """Creates a QuantizedVariable that wraps another variable.
-
-    This typically just returns `QuantizedVariable(variable)`. But, if the variable
-    is a DistributedVariable or one of its subclasses, we instead dynamically
-    create a class that subclasses from both QuantizedVariable and
-    variable.__class__. This is so the returned variable will still pass
-    `isinstance(variable, variable.__class__)`, which is required for
-    DistributedVariables and its subclasses to work properly.
-
-    # Arguments
-    variable: A floating-point resource variable to wrap.
-    quantizer: An optional quantizer to transform the floating-point variable to a
-        fake quantized variable.
-    precision: An optional integer defining the precision of the quantized variable.
-        If `None`, `quantizer.precision` is used.
-
-    # Returns
-    A QuantizedVariable that wraps the variable.
-    """
-    if not isinstance(variable, DistributedVariable):  # type: ignore
-        return QuantizedVariable(variable, quantizer, precision)
-
-    class QuantizedDistributedVariable(QuantizedVariable, variable.__class__):
-        """A QuantizedVariable that also subclasses from DistributedVariable."""
-
-        def get(self, *args, **kwargs):
-            # For some reason this is needed to make unit `x + x` pass on TF 1.14
-            return self._quantize(self.latent_variable.get(*args, **kwargs))
-
-    return QuantizedDistributedVariable(variable, quantizer, precision)
-
-
-def _maybe_wrap(variable, quantizer, precision, wrap=True):
-    """Creates an QuantizedVariable that wraps another variable if applicable.
-
-    This function is used to wrap the return value of QuantizedVariable.assign.
-    Unfortunately MirroredVariable.assign will (incorrectly) return a Mirrored
-    value instead of a MirroredVariable. So we cannot properly wrap it in an
-    AutoCastVariable. We return the original variable in that case.
-
-    # Arguments
-    variable: A tf.Variable or op.
-    quantizer: An optional quantizer to transform the floating-point variable to a
-        fake quantized variable.
-    precision: An optional integer defining the precision of the quantized variable.
-        If `None`, `quantizer.precision` is used.
-    wrap: A boolean to define whether to wrap the variable in an QuantizedVariable.
-
-    # Returns
-    An QuantizedVariable if wrap is True and variable is a resource variable.
-    """
-    if wrap and resource_variable_ops.is_resource_variable(variable):
-        return create_quantized_variable(variable, quantizer, precision)
-    return variable


Test patch:
diff --git a/larq/quantized_variable_test.py b/larq/quantized_variable_test.py
--- a/larq/quantized_variable_test.py
+++ b/larq/quantized_variable_test.py
@@ -5,7 +5,7 @@
 from tensorflow.python.distribute.values import DistributedVariable
 
 from larq import quantized_scope
-from larq.quantized_variable import QuantizedVariable, create_quantized_variable
+from larq.quantized_variable import QuantizedVariable
 from larq.testing_utils import evaluate
 
 
@@ -15,14 +15,14 @@ def get_var(val, dtype=None, name=None):
 
 def test_inheritance(distribute_scope):
     variable = get_var(3.0)
-    quantized_variable = create_quantized_variable(variable)
+    quantized_variable = QuantizedVariable.from_variable(variable)
     assert isinstance(quantized_variable, QuantizedVariable)
     assert isinstance(quantized_variable, tf.Variable)
     assert isinstance(quantized_variable, DistributedVariable) is distribute_scope  # type: ignore
 
 
 def test_read(distribute_scope, eager_and_graph_mode):
-    x = create_quantized_variable(get_var(3.5), quantizer=lambda x: 2 * x)
+    x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     evaluate(x.initializer)
 
     assert evaluate(x) == 3.5
@@ -38,7 +38,7 @@ def test_read(distribute_scope, eager_and_graph_mode):
 
 
 def test_sparse_reads(eager_and_graph_mode):
-    x = create_quantized_variable(get_var([1.0, 2.0]), quantizer=lambda x: 2 * x)
+    x = QuantizedVariable.from_variable(get_var([1.0, 2.0]), quantizer=lambda x: 2 * x)
     evaluate(x.initializer)
 
     assert evaluate(x.sparse_read([0])) == 1
@@ -49,7 +49,7 @@ def test_sparse_reads(eager_and_graph_mode):
 
 
 def test_read_nested_scopes(distribute_scope, eager_and_graph_mode):
-    x = create_quantized_variable(get_var(3.5), quantizer=lambda x: 2 * x)
+    x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     evaluate(x.initializer)
     with quantized_scope.scope(True):
         assert evaluate(x.read_value()) == 7
@@ -59,7 +59,7 @@ def test_read_nested_scopes(distribute_scope, eager_and_graph_mode):
 
 
 def test_method_delegations(distribute_scope, eager_and_graph_mode):
-    x = create_quantized_variable(get_var(3.5), quantizer=lambda x: 2 * x)
+    x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     with quantized_scope.scope(True):
         evaluate(x.initializer)
         assert evaluate(x.value()) == 7
@@ -91,7 +91,7 @@ def test_method_delegations(distribute_scope, eager_and_graph_mode):
 
 
 def test_scatter_method_delegations(eager_and_graph_mode):
-    x = create_quantized_variable(get_var([3.5, 4]), quantizer=lambda x: 2 * x)
+    x = QuantizedVariable.from_variable(get_var([3.5, 4]), quantizer=lambda x: 2 * x)
     evaluate(x.initializer)
     with quantized_scope.scope(True):
         assert_array_equal(evaluate(x.value()), [7, 8])
@@ -120,9 +120,9 @@ def slices(val, index):
 
 def test_overloads(quantized, distribute_scope, eager_and_graph_mode):
     if quantized:
-        x = create_quantized_variable(get_var(3.5), quantizer=lambda x: 2 * x)
+        x = QuantizedVariable.from_variable(get_var(3.5), quantizer=lambda x: 2 * x)
     else:
-        x = create_quantized_variable(get_var(7.0))
+        x = QuantizedVariable.from_variable(get_var(7.0))
     evaluate(x.initializer)
     assert_almost_equal(8, evaluate(x + 1))
     assert_almost_equal(10, evaluate(3 + x))
@@ -155,11 +155,11 @@ def test_overloads(quantized, distribute_scope, eager_and_graph_mode):
 
 def test_tensor_equality(quantized, eager_mode):
     if quantized:
-        x = create_quantized_variable(
+        x = QuantizedVariable.from_variable(
             get_var([3.5, 4.0, 4.5]), quantizer=lambda x: 2 * x
         )
     else:
-        x = create_quantized_variable(get_var([7.0, 8.0, 9.0]))
+        x = QuantizedVariable.from_variable(get_var([7.0, 8.0, 9.0]))
     evaluate(x.initializer)
     assert_array_equal(evaluate(x), [7.0, 8.0, 9.0])
     if version.parse(tf.__version__) >= version.parse("2"):
@@ -168,7 +168,9 @@ def test_tensor_equality(quantized, eager_mode):
 
 
 def test_assign(quantized, distribute_scope, eager_and_graph_mode):
-    x = create_quantized_variable(get_var(0.0, tf.float64), quantizer=lambda x: 2 * x)
+    x = QuantizedVariable.from_variable(
+        get_var(0.0, tf.float64), quantizer=lambda x: 2 * x
+    )
     evaluate(x.initializer)
 
     latent_value = 3.14
@@ -214,7 +216,7 @@ def test_assign(quantized, distribute_scope, eager_and_graph_mode):
 
 
 def test_checkpoint(tmp_path, eager_and_graph_mode):
-    x = create_quantized_variable(get_var(0.0), quantizer=lambda x: 2 * x)
+    x = QuantizedVariable.from_variable(get_var(0.0), quantizer=lambda x: 2 * x)
     evaluate(x.initializer)
     evaluate(x.assign(123.0))
 
@@ -230,11 +232,11 @@ def test_checkpoint(tmp_path, eager_and_graph_mode):
 
 def test_invalid_wrapped_usage(distribute_scope):
     with pytest.raises(ValueError, match="`variable` must be of type"):
-        create_quantized_variable(tf.constant([1.0]))
+        QuantizedVariable.from_variable(tf.constant([1.0]))
     with pytest.raises(ValueError, match="`quantizer` must be `callable` or `None`"):
-        create_quantized_variable(get_var([1.0]), 1)
+        QuantizedVariable.from_variable(get_var([1.0]), 1)
     with pytest.raises(ValueError, match="`precision` must be of type `int` or `None`"):
-        create_quantized_variable(get_var([1.0]), precision=1.0)
+        QuantizedVariable.from_variable(get_var([1.0]), precision=1.0)
 
 
 def test_repr(snapshot, eager_and_graph_mode):
@@ -244,14 +246,18 @@ class Quantizer:
         def __call__(self, x):
             return x
 
-    snapshot.assert_match(repr(create_quantized_variable(x, quantizer=lambda x: 2 * x)))
-    snapshot.assert_match(repr(create_quantized_variable(x, quantizer=Quantizer())))
-    snapshot.assert_match(repr(create_quantized_variable(x, precision=1)))
+    snapshot.assert_match(
+        repr(QuantizedVariable.from_variable(x, quantizer=lambda x: 2 * x))
+    )
+    snapshot.assert_match(
+        repr(QuantizedVariable.from_variable(x, quantizer=Quantizer()))
+    )
+    snapshot.assert_match(repr(QuantizedVariable.from_variable(x, precision=1)))
 
 
 @pytest.mark.parametrize("should_quantize", [True, False])
 def test_optimizer(eager_mode, should_quantize):
-    x = create_quantized_variable(get_var(1.0), quantizer=lambda x: -x)
+    x = QuantizedVariable.from_variable(get_var(1.0), quantizer=lambda x: -x)
     opt = tf.keras.optimizers.SGD(1.0)
 
     def loss():