Source code for gpflow.base

# Copyright 2017-2020 The GPflow Contributors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

from enum import Enum
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple, Union

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from typing_extensions import Final

from .config import default_float, default_summary_fmt
from .experimental.check_shapes import check_shapes
from .type_flags import GENERIC_NP_ARRAYS, NP_TYPE_CHECKING

if TYPE_CHECKING:  # pragma: no cover
    from IPython.lib import pretty

DType = Union[np.dtype, tf.DType]


if TYPE_CHECKING and (not NP_TYPE_CHECKING):  # pragma: no cover
    AnyNDArray = Any
else:
    if GENERIC_NP_ARRAYS:
        # It would be nice to use something more interesting than `Any` here, but it looks like
        # the infrastructure in the rest of the ecosystem isn't really set up for this
        # yet. Maybe when we get Python 3.11?
        AnyNDArray = np.ndarray[Any, Any]  # type: ignore[misc]
    else:
        AnyNDArray = Union[np.ndarray]  # type: ignore[misc]

VariableData = Union[List[Any], Tuple[Any], AnyNDArray, int, float]  # deprecated
Transform = Union[tfp.bijectors.Bijector]
Prior = Union[tfp.distributions.Distribution]


# We've left this as object until we've tested the performance consequences of using the full set
# (np.ndarray, tf.Tensor, tf.Variable, Parameter), see https://github.com/GPflow/GPflow/issues/1434
TensorLike: Final[Tuple[type, ...]] = (object,)
"""
:var TensorLike: Collection of tensor-like types for registering implementations with
    `multipledispatch` dispatchers.
"""


_NativeScalar = Union[int, float]
_Array = Sequence[Any]  # a nested array of int, float, bool etc. kept simple for readability

MeanAndVariance = Tuple[tf.Tensor, tf.Tensor]
SamplesMeanAndVariance = Tuple[tf.Tensor, tf.Tensor, tf.Tensor]


def _IS_PARAMETER(o: object) -> bool:
    return isinstance(o, Parameter)


def _IS_TRAINABLE_PARAMETER(o: object) -> bool:
    return isinstance(o, Parameter) and o.trainable


[docs]class Module(tf.Module): @property def parameters(self) -> Tuple["Parameter", ...]: return tuple(self._flatten(predicate=_IS_PARAMETER)) @property def trainable_parameters(self) -> Tuple["Parameter", ...]: return tuple(self._flatten(predicate=_IS_TRAINABLE_PARAMETER)) def _representation_table(self, object_name: str, tablefmt: Optional[str]) -> str: from .utilities import leaf_components, tabulate_module_summary repr_components = [object_name] if leaf_components(self): repr_components.append(tabulate_module_summary(self, tablefmt=tablefmt)) return "\n".join(repr_components) def _repr_html_(self) -> str: """ Nice representation of GPflow objects in IPython/Jupyter notebooks """ from html import escape return self._representation_table(escape(repr(self)), "html") def _repr_pretty_(self, p: "pretty.RepresentationPrinter", cycle: bool) -> None: """ Nice representation of GPflow objects in the IPython shell """ repr_str = self._representation_table(repr(self), default_summary_fmt()) p.text(repr_str)
[docs]class PriorOn(Enum): CONSTRAINED = "constrained" UNCONSTRAINED = "unconstrained"
[docs]class Parameter(tfp.util.TransformedVariable): def __init__( self, value: "TensorData", *, transform: Optional[Transform] = None, prior: Optional[Prior] = None, prior_on: Optional[Union[str, PriorOn]] = None, trainable: Optional[bool] = None, dtype: Optional[DType] = None, name: Optional[str] = None, unconstrained_shape: Optional[Sequence[Optional[int]]] = None, constrained_shape: Optional[Sequence[Optional[int]]] = None, shape: Optional[Sequence[Optional[int]]] = None, ): """A parameter retains both constrained and unconstrained representations. If no transform is provided, these two values will be the same. It is often challenging to operate with unconstrained parameters. For example, a variance cannot be negative, 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. :param unconstrained_shape: Declare the shape of the unconstrained / pre-transformed values. Useful for setting dynamic shapes. :param constrained_shape: Declare the shape of the constrained / transformed values. Useful for setting dynamic shapes. :param shape: Convenience shortcut for setting both `unconstrained_shape` and `constrained_shape` to the same value. """ if transform: name = name or transform.name if isinstance(value, Parameter): transform = transform or value.transform prior = prior or value.prior prior_on = prior_on or value.prior_on name = name or value.bijector.name trainable = value.trainable if trainable is None else trainable if dtype: tensor_value: TensorType = _cast_to_dtype(value, dtype) else: tensor_value = value else: if transform is None: transform = tfp.bijectors.Identity() prior_on = prior_on if prior_on else PriorOn.CONSTRAINED trainable = trainable if trainable is not None else True tensor_value = _cast_to_dtype(value, dtype) _validate_unconstrained_value(tensor_value, transform, dtype) if shape is not None: assert unconstrained_shape is None, "Cannot set both `shape` and `unconstrained_shape`." assert constrained_shape is None, "Cannot set both `shape` and `constrained_shape`." unconstrained_shape = shape constrained_shape = shape super().__init__( tensor_value, transform, dtype=tensor_value.dtype, trainable=trainable, name=name, shape=unconstrained_shape, ) # TransformedVariable.__init__ doesn't allow us to pass an unconstrained / pre-transformed # shape, so we manually override it. if constrained_shape is not None: self._shape = tf.TensorShape(constrained_shape) self.prior: Optional[Prior] = prior self.prior_on = prior_on # type: ignore[assignment] # see https://github.com/python/mypy/issues/3004
[docs] @check_shapes("return: []") 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 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_variable log_p = tf.reduce_sum(self.prior.log_prob(x)) if self.transform is not None: # need to include log|Jacobian| to account for coordinate transform log_det_jacobian = self.transform.inverse_log_det_jacobian(y, y.shape.ndims) log_p += tf.reduce_sum(log_det_jacobian) return log_p
@property def prior_on(self) -> PriorOn: return self._prior_on @prior_on.setter def prior_on(self, value: Union[str, PriorOn]) -> None: self._prior_on = PriorOn(value) @property def unconstrained_variable(self) -> tf.Variable: return self._pretransformed_input @property def transform(self) -> Optional[Transform]: return self.bijector @property def trainable(self) -> bool: """ `True` if this instance is trainable, else `False`. This attribute cannot be set directly. Use :func:`gpflow.set_trainable`. """ return self.unconstrained_variable.trainable # type: ignore[no-any-return]
[docs] def assign( self, value: "TensorData", use_locking: bool = False, name: Optional[str] = None, read_value: bool = True, ) -> tf.Tensor: """ Assigns constrained `value` to the unconstrained parameter's variable. It passes constrained value through parameter's transform first. Example:: a = Parameter(2.0, transform=tfp.bijectors.Softplus()) b = Parameter(3.0) a.assign(4.0) # `a` parameter to `2.0` value. a.assign(tf.constant(5.0)) # `a` parameter to `5.0` value. a.assign(b) # `a` parameter to constrained value of `b`. :param value: Constrained tensor-like value. :param use_locking: If `True`, use locking during the assignment. :param name: The name of the operation to be created. :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 = _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 )
# These types are defined after "Parameter" to avoid forward references that breaks our # documentation build: TensorType = Union[AnyNDArray, tf.Tensor, tf.Variable, Parameter] """ Type alias for tensor-like types that are supported by most TensorFlow and GPflow operations. NOTE: Union types like this do not work with the `register` method of `multipledispatch`'s `Dispatcher` class. Instead use `TensorLike`. """ TensorData = Union[_NativeScalar, _Array, TensorType] InputData = Union[TensorType] OutputData = Union[TensorType] RegressionData = Tuple[InputData, OutputData] def _cast_to_dtype( value: TensorData, dtype: Optional[DType] = None ) -> Union[tf.Tensor, tf.Variable]: if dtype is None: dtype = default_float() if tf.is_tensor(value): # NOTE(awav) TF2.2 resolves issue with cast. # From TF2.2, `tf.cast` can be used alone instead of this auxiliary function. # workaround for https://github.com/tensorflow/tensorflow/issues/35938 return tf.cast(value, dtype) else: 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) if unconstrained_value.dtype.is_integer: return unconstrained_value 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) return value def _to_unconstrained(value: TensorType, transform: Optional[Transform]) -> TensorType: if transform is not None: return transform.inverse(value) return value