Update parameter to be savable

doc/source/notebooks/theory/FITCvsVFE.py

@@ -103,9 +103,9 @@ def repeatMinimization(model, Xtest, Ytest):
 
 def trainSparseModel(Xtrain, Ytrain, exact_model, isFITC, Xtest, Ytest):
     sparse_model = getSparseModel(Xtrain, Ytrain, isFITC)
-    sparse_model.likelihood.variance = exact_model.likelihood.variance.read_value().copy()
-    sparse_model.kern.lengthscales = exact_model.kern.lengthscales.read_value().copy()
-    sparse_model.kern.variance = exact_model.kern.variance.read_value().copy()
+    sparse_model.likelihood.variance = exact_model.likelihood.variance.numpy()
+    sparse_model.kern.lengthscales = exact_model.kern.lengthscales.numpy()
+    sparse_model.kern.variance = exact_model.kern.variance.numpy()
     return sparse_model, repeatMinimization(sparse_model, Xtest, Ytest)
 
 

doc/source/notebooks/understanding/models.pct.py

@@ -119,10 +119,11 @@
 p.transform.inverse(p)
 
 # %% [markdown]
-# You can also change the `transform` attribute in place:
+# To replace the `transform` of a parameter you need to recreate the parameter with updated transform:
 
 # %%
-m.kernel.kernels[0].variance.transform = tfp.bijectors.Exp()
+p = m.kernel.kernels[0].variance
+m.kernel.kernels[0].variance = gpflow.Parameter(p.numpy(), transform=tfp.bijectors.Exp())
 
 # %%
 print_summary(m, fmt="notebook")
@@ -281,7 +282,7 @@ def maximum_log_likelihood_objective(self):
 # %%
 xx, yy = np.mgrid[-4:4:200j, -4:4:200j]
 X_test = np.vstack([xx.flatten(), yy.flatten()]).T
-f_test = np.dot(X_test, m.W.read_value()) + m.b.read_value()
+f_test = np.dot(X_test, m.W.numpy()) + m.b.numpy()
 p_test = np.exp(f_test)
 p_test /= p_test.sum(1)[:, None]
 

gpflow/base.py

@@ -84,7 +84,7 @@ class PriorOn(Enum):
     UNCONSTRAINED = "unconstrained"
 
 
-class Parameter(tf.Module):
+class Parameter(tfp.util.TransformedVariable):
     def __init__(
         self,
         value: TensorData,
@@ -103,35 +103,31 @@ def __init__(
         therefore we need a positive constraint and it is natural to use constrained values.
         A prior can be imposed either on the constrained version (default) or on the unconstrained version of the parameter.
         """
-        super().__init__()
+        if transform is None:
+            transform = tfp.bijectors.Identity()
+
+        value = _cast_to_dtype(value, dtype)
+        _validate_unconstrained_value(value, transform, dtype)
+        super().__init__(value, transform, dtype=value.dtype, trainable=trainable, name=name)
 
-        self._transform = transform
         self.prior = prior
         self.prior_on = prior_on  # type: ignore  # see https://github.com/python/mypy/issues/3004
 
-        if isinstance(value, tf.Variable):
-            self._unconstrained = value
-        else:
-            unconstrained_value = self.validate_unconstrained_value(value, dtype)
-            self._unconstrained = tf.Variable(
-                unconstrained_value, dtype=dtype, name=name, trainable=trainable
-            )
-
     def log_prior_density(self) -> tf.Tensor:
         """ Log of the prior probability density of the constrained variable. """
 
         if self.prior is None:
             return tf.convert_to_tensor(0.0, dtype=self.dtype)
 
-        y = self.read_value()
+        y = self
 
         if self.prior_on == PriorOn.CONSTRAINED:
             # evaluation is in same space as prior
             return tf.reduce_sum(self.prior.log_prob(y))
 
         else:
             # prior on unconstrained, but evaluating log-prior in constrained space
-            x = self._unconstrained
+            x = self.unconstrained_variable
             log_p = tf.reduce_sum(self.prior.log_prob(x))
 
             if self.transform is not None:
@@ -149,31 +145,13 @@ def prior_on(self) -> PriorOn:
     def prior_on(self, value: Union[str, PriorOn]) -> None:
         self._prior_on = PriorOn(value)
 
-    def value(self) -> tf.Tensor:
-        return _to_constrained(self._unconstrained.value(), self.transform)  # type: ignore  # assumes _to_constrained returns a tf.Tensor
-
-    def read_value(self) -> tf.Tensor:
-        return _to_constrained(self._unconstrained.read_value(), self.transform)  # type: ignore  # assumes _to_constrained returns a tf.Tensor
-
-    def experimental_ref(self) -> "Parameter":
-        return self
-
-    def deref(self) -> "Parameter":
-        return self
-
     @property
     def unconstrained_variable(self) -> tf.Variable:
-        return self._unconstrained
+        return self._pretransformed_input
 
     @property
     def transform(self) -> Optional[Transform]:
-        return self._transform
-
-    @transform.setter
-    def transform(self, new_transform: Optional[Transform]) -> None:
-        constrained_value = self.read_value()
-        self._transform = new_transform
-        self.assign(constrained_value)
+        return self.bijector
 
     @property
     def trainable(self) -> bool:
@@ -182,21 +160,7 @@ def trainable(self) -> bool:
 
         This attribute cannot be set directly. Use :func:`gpflow.set_trainable`.
         """
-        return self._unconstrained.trainable
-
-    @property
-    def initial_value(self) -> tf.Tensor:
-        return self._unconstrained.initial_value
-
-    def validate_unconstrained_value(self, value: TensorData, dtype: DType) -> tf.Tensor:
-        value = _cast_to_dtype(value, dtype)
-        unconstrained_value = _to_unconstrained(value, self.transform)
-        message = (
-            "gpflow.Parameter: the value to be assigned is incompatible with this parameter's "
-            "transform (the corresponding unconstrained value has NaN or Inf) and hence cannot be "
-            "assigned."
-        )
-        return tf.debugging.assert_all_finite(unconstrained_value, message=message)
+        return self.unconstrained_variable.trainable
 
     def assign(
         self,
@@ -225,125 +189,11 @@ def assign(
         :param read_value: if True, will return something which evaluates to the new
             value of the variable; if False will return the assign op.
         """
-        unconstrained_value = self.validate_unconstrained_value(value, self.dtype)
-        return self._unconstrained.assign(
+        unconstrained_value = _validate_unconstrained_value(value, self.transform, self.dtype)
+        return self.unconstrained_variable.assign(
             unconstrained_value, use_locking=use_locking, name=name, read_value=read_value
         )
 
-    @property
-    def is_tensor_like(self) -> bool:
-        """
-        This method means that TensorFlow's `tensor_util.is_tensor` function
-        will return `True`
-        """
-        return True
-
-    @property
-    def name(self) -> str:
-        return self._unconstrained.name
-
-    @property
-    def initializer(self):  # type unknown
-        return self._unconstrained.initializer
-
-    @property
-    def device(self) -> Optional[str]:
-        return self._unconstrained.device
-
-    @property
-    def dtype(self) -> tf.DType:
-        return self._unconstrained.dtype
-
-    @property
-    def op(self) -> tf.Operation:
-        return self._unconstrained.op
-
-    @property
-    def shape(self) -> tf.TensorShape:
-        if self.transform is not None:
-            return self.transform.forward_event_shape(self._unconstrained.shape)
-        return self._unconstrained.shape
-
-    def numpy(self) -> np.ndarray:
-        return self.read_value().numpy()
-
-    def get_shape(self) -> tf.TensorShape:
-        return self.shape
-
-    def _should_act_as_resource_variable(self):  # type unknown
-        # needed so that Parameters are correctly identified by TensorFlow's
-        # is_resource_variable() in resource_variable_ops.py
-        pass  # only checked by TensorFlow using hasattr()
-
-    @property
-    def handle(self):  # type unknown
-        return self._unconstrained.handle
-
-    def __repr__(self) -> str:
-        unconstrained = self.unconstrained_variable
-        constrained = self.read_value()
-        if tf.executing_eagerly():
-            info = (
-                f"unconstrained-shape={unconstrained.shape} "
-                f"unconstrained-value={unconstrained.numpy()} "
-                f"constrained-shape={constrained.shape} "
-                f"constrained-value={constrained.numpy()}"
-            )
-        else:
-            if unconstrained.shape == constrained.shape:
-                info = f"shape={constrained.shape}"
-            else:
-                info = (
-                    f"unconstrained-shape={unconstrained.shape} "
-                    f"constrained-shape={constrained.shape}"
-                )
-
-        return f"<gpflow.Parameter {self.name!r} dtype={self.dtype.name} {info}>"
-
-    # Below
-    # TensorFlow copy-paste code to make variable-like object to work
-
-    @classmethod
-    def _OverloadAllOperators(cls):  # pylint: disable=invalid-name
-        """Register overloads for all operators."""
-        for operator in tf.Tensor.OVERLOADABLE_OPERATORS:
-            cls._OverloadOperator(operator)
-        # For slicing, bind getitem differently than a tensor (use SliceHelperVar
-        # instead)
-        # pylint: disable=protected-access
-        setattr(cls, "__getitem__", array_ops._SliceHelperVar)
-
-    @classmethod
-    def _OverloadOperator(cls, operator):  # pylint: disable=invalid-name
-        """Defer an operator overload to `ops.Tensor`.
-
-        We pull the operator out of ops.Tensor dynamically to avoid ordering issues.
-
-        Args:
-            operator: string. The operator name.
-        """
-        tensor_oper = getattr(tf.Tensor, operator)
-
-        def _run_op(a, *args, **kwargs):
-            # pylint: disable=protected-access
-            return tensor_oper(a.read_value(), *args, **kwargs)
-
-        functools.update_wrapper(_run_op, tensor_oper)
-        setattr(cls, operator, _run_op)
-
-    # NOTE(mrry): This enables the Variable's overloaded "right" binary
-    # operators to run when the left operand is an ndarray, because it
-    # accords the Variable class higher priority than an ndarray, or a
-    # numpy matrix.
-    # TODO(mrry): Convert this to using numpy's __numpy_ufunc__
-    # mechanism, which allows more control over how Variables interact
-    # with ndarrays.
-    __array_priority__ = 100
-
-
-Parameter._OverloadAllOperators()
-tf.register_tensor_conversion_function(Parameter, lambda x, *args, **kwds: x.read_value())
-
 
 def _cast_to_dtype(
     value: TensorData, dtype: Optional[DType] = None
@@ -360,6 +210,19 @@ def _cast_to_dtype(
         return tf.convert_to_tensor(value, dtype=dtype)
 
 
+def _validate_unconstrained_value(
+    value: TensorData, transform: tfp.bijectors.Bijector, dtype: DType
+) -> tf.Tensor:
+    value = _cast_to_dtype(value, dtype)
+    unconstrained_value = _to_unconstrained(value, transform)
+    message = (
+        "gpflow.Parameter: the value to be assigned is incompatible with this parameter's "
+        "transform (the corresponding unconstrained value has NaN or Inf) and hence cannot be "
+        "assigned."
+    )
+    return tf.debugging.assert_all_finite(unconstrained_value, message=message)
+
+
 def _to_constrained(value: TensorType, transform: Optional[Transform]) -> TensorType:
     if transform is not None:
         return transform.forward(value)

gpflow/optimizers/natgrad.py

@@ -195,15 +195,14 @@ def _natgrad_steps(
         :param parameters: List of tuples (q_mu, q_sqrt, xi_transform)
         """
         q_mus, q_sqrts, xis = zip(*parameters)
-        unconstrained_variables = [
-            p.unconstrained_variable for params in (q_mus, q_sqrts) for p in params
-        ]
+        q_mu_vars = [p.unconstrained_variable for p in q_mus]
+        q_sqrt_vars = [p.unconstrained_variable for p in q_sqrts]
 
         with tf.GradientTape(watch_accessed_variables=False) as tape:
-            tape.watch(unconstrained_variables)
+            tape.watch(q_mu_vars + q_sqrt_vars)
             loss = loss_fn()
 
-        q_mu_grads, q_sqrt_grads = tape.gradient(loss, [q_mus, q_sqrts])
+        q_mu_grads, q_sqrt_grads = tape.gradient(loss, [q_mu_vars, q_sqrt_vars])
         # NOTE that these are the gradients in *unconstrained* space
 
         with tf.name_scope(f"{self._name}/natural_gradient_steps"):

gpflow/utilities/utilities.py

@@ -190,17 +190,22 @@ def leaf_components(input: tf.Module):
 def _merge_leaf_components(
     input: Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]
 ) -> Dict[str, Union[tf.Variable, tf.Tensor, Parameter]]:
-    input_values = set([value.experimental_ref() for value in input.values()])
+
+    ref_fn = lambda x: (x if isinstance(x, Parameter) else x.ref())
+    deref_fn = lambda x: (x if isinstance(x, Parameter) else x.deref())
+
+    input_values = set([ref_fn(value) for value in input.values()])
     if len(input_values) == len(input):
         return input
+
     tmp_dict = dict()  # Type: Dict[ref, str]
-    for key, variable in input.items():
-        ref = variable.experimental_ref()
+    for key, value in input.items():
+        ref = ref_fn(value)
         if ref in tmp_dict:
             tmp_dict[ref] = f"{tmp_dict[ref]}\n{key}"
         else:
             tmp_dict[ref] = key
-    return {key: ref.deref() for ref, key in tmp_dict.items()}
+    return {key: deref_fn(ref) for ref, key in tmp_dict.items()}
 
 
 def _get_leaf_components(input_module: tf.Module):

tests/gpflow/expectations/test_expectations.py

@@ -214,7 +214,7 @@ def test_RBF_eKzxKxz_gradient_notNaN():
 
     with tf.GradientTape() as tape:
         ekz = expectation(p, (kernel, z), (kernel, z))
-    grad = tape.gradient(ekz, kernel.lengthscales)
+    grad = tape.gradient(ekz, kernel.lengthscales.unconstrained_variable)
     assert grad is not None and not np.isnan(grad)