bject for storing global GPflow settings
    """

    int: type = field(default_factory=_default_int_factory)
    """Integer data type, int32 or int64."""

    float: type = field(default_factory=_default_float_factory)
    """Float data type, float32 or float64"""

    jitter: Float = field(default_factory=_default_jitter_factory)
    """
    Jitter value. Mainly used for for making badly conditioned matrices more stable.

    Default value is `1e-6`.
    """

    positive_bijector: str = field(default_factory=_default_positive_bijector_factory)
    """
    Method for positive bijector, either "softplus" or "exp".

    Default is "softplus".
    """

    positive_minimum: Float = field(default_factory=_default_positive_minimum_factory)
    """Lower bound for the positive transformation."""

    summary_fmt: Optional[str] = field(default_factory=_default_summary_fmt_factory)
    """Summary format for module printing."""


def config() -> Config:
    """Returns current active config."""
    assert __config is not None, "__config is None. This should never happen."
    return __config


def default_int() -> type:
    """Returns default integer type"""
    return config().int


def default_float() -> type:
    """Returns default float type"""
    return config().float


def default_jitter() -> float:
    """
    The jitter is a constant that GPflow adds to the diagonal of matrices
    to achieve numerical stability of the system when the condition number
    of the associated matrices is large, and therefore the matrices nearly singular.
    """
    return config().jitter


def default_positive_bijector() -> str:
    """Type of bijector used for positive constraints: exp or softplus."""
    return config().positive_bijector


def default_positive_minimum() -> float:
    """Shift constant that GPflow adds to all positive constraints."""
    return config().positive_minimum


def default_summary_fmt() -> Optional[str]:
    """Summary printing format as understood by :mod:`tabulate` or a special case "notebook"."""
    return config().summary_fmt


def set_config(new_config: Config) -> None:
    """Update GPflow config with new settings from `new_config`."""
    global __config
    __config = new_config


def set_default_int(value_type: type) -> None:
    """
    Sets default integer type. Available options are ``np.int16``, ``np.int32``,
    or ``np.int64``.
    """
    try:
        tf_dtype = tf.as_dtype(value_type)  # Test that it's a tensorflow-valid dtype
    except TypeError:
        raise TypeError(f"{value_type} is not a valid tf or np dtype")

    if not tf_dtype.is_integer:
        raise TypeError(f"{value_type} is not an integer dtype")

    set_config(replace(config(), int=tf_dtype.as_numpy_dtype))


def set_default_float(value_type: type) -> None:
    """
    Sets default float type. Available options are `np.float16`, `np.float32`,
    or `np.float64`.
    """
    try:
        tf_dtype = tf.as_dtype(value_type)  # Test that it's a tensorflow-valid dtype
    except TypeError:
        raise TypeError(f"{value_type} is not a valid tf or np dtype")

    if not tf_dtype.is_floating:
        raise TypeError(f"{value_type} is not a float dtype")

    set_config(replace(config(), float=tf_dtype.as_numpy_dtype))


def set_default_jitter(value: float) -> None:
    """
    Sets constant jitter value.
    The jitter is a constant that GPflow adds to the diagonal of matrices
    to achieve numerical stability of the system when the condition number
    of the associated matrices is large, and therefore the matrices nearly singular.
    """
    if not (
        isinstance(value, (tf.Tensor, np.ndarray)) and len(value.shape) == 0
    ) and not isinstance(value, float):
        raise TypeError("Expected float32 or float64 scalar value")

    if value < 0:
        raise ValueError("Jitter must be non-negative")

    set_config(replace(config(), jitter=value))


def set_default_positive_bijector(value: str) -> None:
    """
    Sets positive bijector type.
    There are currently two options implemented: "exp" and "softplus".
    """
    type_map = positive_bijector_type_map()
    if isinstance(value, str):
        value = value.lower()
    if value not in type_map:
        raise ValueError(f"`{value}` not in set of valid bijectors: {sorted(type_map)}")

    set_config(replace(config(), positive_bijector=value))


def set_default_positive_minimum(value: float) -> None:
    """Sets shift constant for positive transformation."""
    if not (
        isinstance(value, (tf.Tensor, np.ndarray)) and len(value.shape) == 0
    ) and not isinstance(value, float):
        raise TypeError("Expected float32 or float64 scalar value")

    if value < 0:
        raise ValueError("Positive minimum must be non-negative")

    set_config(replace(config(), positive_minimum=value))


def set_default_summary_fmt(value: Optional[str]) -> None:
    formats: List[Optional[str]] = list(tabulate.tabulate_formats)
    formats.extend(["notebook", None])
    if value not in formats:
        raise ValueError(f"Summary does not support '{value}' format")

    set_config(replace(config(), summary_fmt=value))


def positive_bijector_type_map() -> Dict[str, type]:
    return {
        "exp": tfp.bijectors.Exp,
        "softplus": tfp.bijectors.Softplus,
    }


@contextlib.contextmanager
def as_context(temporary_config: Optional[Config] = None) -> Generator[None, None, None]:
    """Ensure that global configs defaults, with a context manager. Useful for testing."""
    current_config = config()
    temporary_config = replace(current_config) if temporary_config is None else temporary_config
    try:
        set_config(temporary_config)
        yield
    finally:
        set_config(current_config)


# Set global config.
set_config(Config())


# gpflow/config/__init__.py

from .__config__ import *
from .__config__ import __doc__  # needs explicit import

__all__ = [
    "Config",
    "as_context",
    "config",
    "default_float",
    "default_int",
    "default_jitter",
    "default_positive_bijector",
    "default_positive_minimum",
    "default_summary_fmt",
    "positive_bijector_type_map",
    "set_config",
    "set_default_float",
    "set_default_int",
    "set_default_jitter",
    "set_default_positive_bijector",
    "set_default_positive_minimum",
    "set_default_summary_fmt",
]


# gpflow/covariances/__init__.py

from . import kufs, kuus, multioutput
from .dispatch import Kuf, Kuu

__all__ = [
    "Kuf",
    "Kuu",
    "dispatch",
    "kufs",
    "kuus",
    "multioutput",
]


# gpflow/covariances/dispatch.py

# 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 ..utilities import Dispatcher

Kuu = Dispatcher("Kuu")
Kuf = Dispatcher("Kuf")


# gpflow/covariances/kufs.py

# 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.

import numpy as np  # pylint: disable=unused-import  # Used by Sphinx to generate documentation.
import tensorflow as tf

from ..base import TensorLike, TensorType
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPatches, InducingPoints, Multiscale
from ..kernels import Convolutional, Kernel, SquaredExponential
from .dispatch import Kuf


@Kuf.register(InducingPoints, Kernel, TensorLike)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "Xnew: [batch..., N, D]",
    "return: [M, batch..., N]",
)
def Kuf_kernel_inducingpoints(
    inducing_variable: InducingPoints, kernel: Kernel, Xnew: TensorType
) -> tf.Tensor:
    return kernel(inducing_variable.Z, Xnew)


@Kuf.register(Multiscale, SquaredExponential, TensorLike)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "Xnew: [batch..., N, D]",
    "return: [M, batch..., N]",
)
def Kuf_sqexp_multiscale(
    inducing_variable: Multiscale, kernel: SquaredExponential, Xnew: TensorType
) -> tf.Tensor:
    Xnew, _ = kernel.slice(Xnew, None)
    Zmu, Zlen = kernel.slice(inducing_variable.Z, inducing_variable.scales)
    idlengthscales = kernel.lengthscales + Zlen
    d = inducing_variable._cust_square_dist(Xnew, Zmu, idlengthscales[None, :, :])
    lengthscales = tf.reduce_prod(kernel.lengthscales / idlengthscales, 1)
    lengthscales = tf.reshape(lengthscales, (1, -1))
    return tf.transpose(kernel.variance * tf.exp(-0.5 * d) * lengthscales)


@Kuf.register(InducingPatches, Convolutional, object)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "Xnew: [batch..., N, D2]",
    "return: [M, batch..., N]",
)
def Kuf_conv_patch(
    inducing_variable: InducingPatches, kernel: Convolutional, Xnew: TensorType
) -> tf.Tensor:
    Xp = kernel.get_patches(Xnew)  # [N, num_patches, patch_len]
    bigKzx = kernel.base_kernel.K(
        inducing_variable.Z, Xp
    )  # [M, N, P] -- thanks to broadcasting of kernels
    Kzx = tf.reduce_sum(bigKzx * kernel.weights if hasattr(kernel, "weights") else bigKzx, [2])
    return Kzx / kernel.num_patches


# gpflow/covariances/kuus.py

# 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.

import tensorflow as tf

from ..config import default_float
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPatches, InducingPoints, Multiscale
from ..kernels import Convolutional, Kernel, SquaredExponential
from .dispatch import Kuu


@Kuu.register(InducingPoints, Kernel)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "return: [M, M]",
)
def Kuu_kernel_inducingpoints(
    inducing_variable: InducingPoints, kernel: Kernel, *, jitter: float = 0.0
) -> tf.Tensor:
    Kzz = kernel(inducing_variable.Z)
    Kzz += jitter * tf.eye(inducing_variable.num_inducing, dtype=Kzz.dtype)
    return Kzz


@Kuu.register(Multiscale, SquaredExponential)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "return: [M, M]",
)
def Kuu_sqexp_multiscale(
    inducing_variable: Multiscale, kernel: SquaredExponential, *, jitter: float = 0.0
) -> tf.Tensor:
    Zmu, Zlen = kernel.slice(inducing_variable.Z, inducing_variable.scales)
    idlengthscales2 = tf.square(kernel.lengthscales + Zlen)
    sc = tf.sqrt(
        idlengthscales2[None, ...] + idlengthscales2[:, None, ...] - kernel.lengthscales ** 2
    )
    d = inducing_variable._cust_square_dist(Zmu, Zmu, sc)
    Kzz = kernel.variance * tf.exp(-d / 2) * tf.reduce_prod(kernel.lengthscales / sc, 2)
    Kzz += jitter * tf.eye(inducing_variable.num_inducing, dtype=Kzz.dtype)
    return Kzz


@Kuu.register(InducingPatches, Convolutional)
@check_shapes(
    "inducing_variable: [M, D, 1]",
    "return: [M, M]",
)
def Kuu_conv_patch(
    inducing_variable: InducingPatches, kernel: Convolutional, jitter: float = 0.0
) -> tf.Tensor:
    return kernel.base_kernel.K(inducing_variable.Z) + jitter * tf.eye(
        inducing_variable.num_inducing, dtype=default_float()
    )


# gpflow/expectations/__init__.py

from . import (
    cross_kernels,
    linears,
    mean_functions,
    misc,
    products,
    quadratures,
    squared_exponentials,
    sums,
)
from .expectations import expectation, quadrature_expectation

__all__ = [
    "cross_kernels",
    "dispatch",
    "expectation",
    "expectations",
    "linears",
    "mean_functions",
    "misc",
    "products",
    "quadrature_expectation",
    "quadratures",
    "squared_exponentials",
    "sums",
]


# gpflow/expectations/cross_kernels.py

# 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 typing import Union, cast

import tensorflow as tf

from .. import kernels
from ..base import TensorType
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian
from . import dispatch
from .expectations import expectation


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian),
    kernels.SquaredExponential,
    InducingPoints,
    kernels.Linear,
    InducingPoints,
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M1, D, P]",
    "feat2: [M2, D, P]",
    "return: [N, M1, M2]",
)
def _expectation_gaussian_sqe_inducingpoints__linear_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    sqexp_kern: kernels.SquaredExponential,
    feat1: InducingPoints,
    lin_kern: kernels.Linear,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - K_lin_{.,.} :: SqExp kernel
        - K_sqexp_{.,.} :: Linear kernel
    Different Z1 and Z2 are handled if p is diagonal and K_lin and K_sqexp have disjoint
    active_dims, in which case the joint expectations simplify into a product of expectations

    :return: NxM1xM2
    """
    if sqexp_kern.on_separate_dims(lin_kern) and isinstance(
        p, DiagonalGaussian
    ):  # no joint expectations required
        eKxz1 = expectation(p, (sqexp_kern, feat1))
        eKxz2 = expectation(p, (lin_kern, feat2))
        return eKxz1[:, :, None] * eKxz2[:, None, :]

    if feat1 != feat2:
        raise NotImplementedError("inducing_variables have to be the same for both kernels.")

    if sqexp_kern.active_dims != lin_kern.active_dims:
        raise NotImplementedError("active_dims have to be the same for both kernels.")

    # use only active dimensions
    Xcov = sqexp_kern.slice_cov(tf.linalg.diag(p.cov) if isinstance(p, DiagonalGaussian) else p.cov)
    Z, Xmu = sqexp_kern.slice(feat1.Z, p.mu)

    N = tf.shape(Xmu)[0]
    D = tf.shape(Xmu)[1]

    def take_with_ard(value: TensorType) -> TensorType:
        if not sqexp_kern.ard:
            return tf.zeros((D,), dtype=value.dtype) + value
        return value

    lin_kern_variances = take_with_ard(lin_kern.variance)
    sqexp_kern_lengthscales = take_with_ard(sqexp_kern.lengthscales)

    chol_L_plus_Xcov = tf.linalg.cholesky(
        tf.linalg.diag(sqexp_kern_lengthscales ** 2) + Xcov
    )  # NxDxD

    Z_transpose = tf.transpose(Z)
    all_diffs = Z_transpose - tf.expand_dims(Xmu, 2)  # NxDxM
    exponent_mahalanobis = tf.linalg.triangular_solve(
        chol_L_plus_Xcov, all_diffs, lower=True
    )  # NxDxM
    exponent_mahalanobis = tf.reduce_sum(tf.square(exponent_mahalanobis), 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    sqrt_det_L = tf.reduce_prod(sqexp_kern_lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N
    eKxz_sqexp = sqexp_kern.variance * (
        determinants[:, None] * exponent_mahalanobis
    )  ## NxM <- End RBF eKxz code

    tiled_Z = tf.tile(tf.expand_dims(Z_transpose, 0), (N, 1, 1))  # NxDxM
    z_L_inv_Xcov = tf.linalg.matmul(
        tiled_Z, Xcov / sqexp_kern_lengthscales[:, None] ** 2.0, transpose_a=True
    )  # NxMxD

    cross_eKzxKxz = tf.linalg.cholesky_solve(
        chol_L_plus_Xcov,
        cast(tf.Tensor, (lin_kern_variances * sqexp_kern_lengthscales ** 2.0))[..., None] * tiled_Z,
    )  # NxDxM

    cross_eKzxKxz = tf.linalg.matmul(
        (z_L_inv_Xcov + Xmu[:, None, :]) * eKxz_sqexp[..., None], cross_eKzxKxz
    )  # NxMxM
    return cross_eKzxKxz


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian),
    kernels.Linear,
    InducingPoints,
    kernels.SquaredExponential,
    InducingPoints,
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M1, D, P]",
    "feat2: [M2, D, P]",
    "return: [N, M1, M2]",
)
def _expectation_gaussian_linear_inducingpoints__sqe_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    lin_kern: kernels.Linear,
    feat1: InducingPoints,
    sqexp_kern: kernels.SquaredExponential,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - K_lin_{.,.} :: Linear kernel
        - K_sqexp_{.,.} :: sqexp kernel
    Different Z1 and Z2 are handled if p is diagonal and K_lin and K_sqexp have disjoint
    active_dims, in which case the joint expectations simplify into a product of expectations

    :return: NxM1xM2
    """
    return tf.linalg.adjoint(expectation(p, (sqexp_kern, feat2), (lin_kern, feat1)))


# gpflow/expectations/dispatch.py

# 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 ..utilities import Dispatcher

expectation = Dispatcher("expectation")
quadrature_expectation = Dispatcher("quadrature_expectation")
variational_expectation = Dispatcher("variational_expectation")


# gpflow/expectations/expectations.py

# Copyright 2018 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 typing import Optional, Tuple, Union, cast

import tensorflow as tf

from ..base import TensorType
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingVariables
from ..kernels import Kernel
from ..mean_functions import MeanFunction
from ..probability_distributions import (
    DiagonalGaussian,
    Gaussian,
    MarkovGaussian,
    ProbabilityDistribution,
)
from . import dispatch

ProbabilityDistributionLike = Union[ProbabilityDistribution, Tuple[TensorType, TensorType]]
"""
Either a prabability distribution, or a tuple of mean, covariance that is turned into an
appropriate Gaussian distribution, depending on the shape of the covariance.
"""

ExpectationObject = Union[Kernel, MeanFunction, None]
PackedExpectationObject = Union[ExpectationObject, Tuple[Kernel, InducingVariables]]


def expectation(
    p: ProbabilityDistributionLike,
    obj1: PackedExpectationObject,
    obj2: PackedExpectationObject = None,
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    Compute the expectation <obj1(x) obj2(x)>_p(x)
    Uses multiple-dispatch to select an analytical implementation,
    if one is available. If not, it falls back to quadrature.

    :type p: (mu, cov) tuple or a `ProbabilityDistribution` object
    :type obj1: kernel, mean function, (kernel, inducing_variable), or None
    :type obj2: kernel, mean function, (kernel, inducing_variable), or None
    :param int nghp: passed to `_quadrature_expectation` to set the number
                     of Gauss-Hermite points used: `num_gauss_hermite_points`
    :return: a 1-D, 2-D, or 3-D tensor containing the expectation

    Allowed combinations

    - Psi statistics:
        >>> eKdiag = expectation(p, kernel)  (N)  # Psi0
        >>> eKxz = expectation(p, (kernel, inducing_variable))  (NxM)  # Psi1
        >>> exKxz = expectation(p, identity_mean, (kernel, inducing_variable))  (NxDxM)
        >>> eKzxKxz = expectation(p, (kernel, inducing_variable), (kernel, inducing_variable))  (NxMxM)  # Psi2

    - kernels and mean functions:
        >>> eKzxMx = expectation(p, (kernel, inducing_variable), mean)  (NxMxQ)
        >>> eMxKxz = expectation(p, mean, (kernel, inducing_variable))  (NxQxM)

    - only mean functions:
        >>> eMx = expectation(p, mean)  (NxQ)
        >>> eM1x_M2x = expectation(p, mean1, mean2)  (NxQ1xQ2)
        .. note:: mean(x) is 1xQ (row vector)

    - different kernels. This occurs, for instance, when we are calculating Psi2 for Sum kernels:
        >>> eK1zxK2xz = expectation(p, (kern1, inducing_variable), (kern2, inducing_variable))  (NxMxM)
    """
    p, obj1, feat1, obj2, feat2 = _init_expectation(p, obj1, obj2)
    try:
        return dispatch.expectation(p, obj1, feat1, obj2, feat2, nghp=nghp)
    except NotImplementedError:
        return dispatch.quadrature_expectation(p, obj1, feat1, obj2, feat2, nghp=nghp)


def quadrature_expectation(
    p: ProbabilityDistributionLike,
    obj1: PackedExpectationObject,
    obj2: PackedExpectationObject = None,
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    Compute the expectation <obj1(x) obj2(x)>_p(x)
    Uses Gauss-Hermite quadrature for approximate integration.

    :type p: (mu, cov) tuple or a `ProbabilityDistribution` object
    :type obj1: kernel, mean function, (kernel, inducing_variable), or None
    :type obj2: kernel, mean function, (kernel, inducing_variable), or None
    :param int num_gauss_hermite_points: passed to `_quadrature_expectation` to set
                                         the number of Gauss-Hermite points used
    :return: a 1-D, 2-D, or 3-D tensor containing the expectation
    """
    print(f"2. p={p}, obj1={obj1}, obj2={obj2}")
    p, obj1, feat1, obj2, feat2 = _init_expectation(p, obj1, obj2)
    return dispatch.quadrature_expectation(p, obj1, feat1, obj2, feat2, nghp=nghp)


@check_shapes(
    "return[0]: [N, D]",
    "return[2]: [M1, D, P]",
    "return[4]: [M2, D, P]",
)
def _init_expectation(
    p: ProbabilityDistributionLike, obj1: PackedExpectationObject, obj2: PackedExpectationObject
) -> Tuple[
    ProbabilityDistribution,
    ExpectationObject,
    Optional[InducingVariables],
    ExpectationObject,
    Optional[InducingVariables],
]:
    if isinstance(p, tuple):
        mu, cov = p
        classes = [DiagonalGaussian, Gaussian, MarkovGaussian]
        p = classes[cov.ndim - 2](*p)  # type: ignore[abstract]

    obj1, feat1 = obj1 if isinstance(obj1, tuple) else (obj1, None)
    obj2, feat2 = obj2 if isinstance(obj2, tuple) else (obj2, None)
    return (
        # type-ignore instead of cast, because it dependes on versions whether a cast is necessary.
        p,  # type: ignore[return-value]
        cast(ExpectationObject, obj1),
        cast(Optional[InducingVariables], feat1),
        cast(ExpectationObject, obj2),
        cast(Optional[InducingVariables], feat2),
    )


# gpflow/expectations/linears.py

# 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 typing import Type, Union

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, kernels.Linear, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_gaussian_linear(
    p: Gaussian, kernel: kernels.Linear, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    <diag(K_{X, X})>_p(X)
        - K_{.,.} :: Linear kernel

    :return: N
    """
    # use only active dimensions
    Xmu, _ = kernel.slice(p.mu, None)
    Xcov = kernel.slice_cov(p.cov)

    return tf.reduce_sum(kernel.variance * (tf.linalg.diag_part(Xcov) + Xmu ** 2), 1)


@dispatch.expectation.register(Gaussian, kernels.Linear, InducingPoints, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_gaussian_linear_inducingpoints(
    p: Gaussian,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    _: None,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    <K_{X, Z}>_p(X)
        - K_{.,.} :: Linear kernel

    :return: NxM
    """
    # use only active dimensions
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)

    return tf.linalg.matmul(Xmu, Z * kernel.variance, transpose_b=True)


@dispatch.expectation.register(Gaussian, kernels.Linear, InducingPoints, mfn.Identity, NoneType)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M, D]",
)
def _expectation_gaussian_linear_inducingpoints__identity(
    p: Gaussian,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    mean: mfn.Identity,
    _: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <K_{Z, x_n} x_n^T>_p(x_n)
        - K_{.,.} :: Linear kernel

    :return: NxMxD
    """
    Xmu, Xcov = p.mu, p.cov

    N = tf.shape(Xmu)[0]
    var_Z = kernel.variance * inducing_variable.Z  # MxD
    tiled_Z = tf.tile(tf.expand_dims(var_Z, 0), (N, 1, 1))  # NxMxD
    return tf.linalg.matmul(tiled_Z, Xcov + (Xmu[..., None] * Xmu[:, None, :]))


@dispatch.expectation.register(
    MarkovGaussian, kernels.Linear, InducingPoints, mfn.Identity, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M, D]",
)
def _expectation_markov_linear_inducingpoints__identity(
    p: MarkovGaussian,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    mean: mfn.Identity,
    _: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <K_{Z, x_n} x_{n+1}^T>_p(x_{n:n+1})
        - K_{.,.} :: Linear kernel
        - p       :: MarkovGaussian distribution (p.cov 2x(N+1)xDxD)

    :return: NxMxD
    """
    Xmu, Xcov = p.mu, p.cov

    N = tf.shape(Xmu)[0] - 1
    var_Z = kernel.variance * inducing_variable.Z  # MxD
    tiled_Z = tf.tile(tf.expand_dims(var_Z, 0), (N, 1, 1))  # NxMxD
    eXX = Xcov[1, :-1] + (Xmu[:-1][..., None] * Xmu[1:][:, None, :])  # NxDxD
    return tf.linalg.matmul(tiled_Z, eXX)


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian), kernels.Linear, InducingPoints, kernels.Linear, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M, D, P]",
    "feat2: [M, D, P]",
    "return: [N, M, M]",
)
def _expectation_gaussian_linear_inducingpoints__linear_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    kern1: kernels.Linear,
    feat1: InducingPoints,
    kern2: kernels.Linear,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - Ka_{.,.}, Kb_{.,.} :: Linear kernels
    Ka and Kb as well as Z1 and Z2 can differ from each other, but this is supported
    only if the Gaussian p is Diagonal (p.cov NxD) and Ka, Kb have disjoint active_dims
    in which case the joint expectations simplify into a product of expectations

    :return: NxMxM
    """
    if kern1.on_separate_dims(kern2) and isinstance(
        p, DiagonalGaussian
    ):  # no joint expectations required
        eKxz1 = expectation(p, (kern1, feat1))
        eKxz2 = expectation(p, (kern2, feat2))
        return eKxz1[:, :, None] * eKxz2[:, None, :]

    if kern1 != kern2 or feat1 != feat2:
        raise NotImplementedError(
            "The expectation over two kernels has only an "
            "analytical implementation if both kernels are equal."
        )

    kernel = kern1
    inducing_variable = feat1

    # use only active dimensions
    Xcov = kernel.slice_cov(tf.linalg.diag(p.cov) if isinstance(p, DiagonalGaussian) else p.cov)
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)

    N = tf.shape(Xmu)[0]
    var_Z = kernel.variance * Z
    tiled_Z = tf.tile(tf.expand_dims(var_Z, 0), (N, 1, 1))  # NxMxD
    XX = Xcov + tf.expand_dims(Xmu, 1) * tf.expand_dims(Xmu, 2)  # NxDxD
    return tf.linalg.matmul(tf.linalg.matmul(tiled_Z, XX), tiled_Z, transpose_b=True)


# gpflow/expectations/mean_functions.py

# 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.

# noqa: F811

from typing import Type, Union

import tensorflow as tf

from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..probability_distributions import Gaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, (mfn.Linear, mfn.Constant), NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q]",
)
def _expectation_gaussian_linear(
    p: Gaussian,
    mean: Union[mfn.Linear, mfn.Constant],
    _: None,
    __: None,
    ___: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    <m(X)>_p(X)
        - m(x) :: Linear, Identity or Constant mean function

    :return: NxQ
    """
    return mean(p.mu)


@dispatch.expectation.register(Gaussian, mfn.Constant, NoneType, mfn.Constant, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_constant__constant(
    p: Gaussian, mean1: mfn.Constant, _: None, mean2: mfn.Constant, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.), m2(.) :: Constant mean functions

    :return: NxQ1xQ2
    """
    return mean1(p.mu)[:, :, None] * mean2(p.mu)[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.Constant, NoneType, mfn.MeanFunction, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_constant__meanfunction(
    p: Gaussian, mean1: mfn.Constant, _: None, mean2: mfn.Constant, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: Constant mean function
        - m2(.) :: General mean function

    :return: NxQ1xQ2
    """
    e_mean2 = expectation(p, mean2)
    return mean1(p.mu)[:, :, None] * e_mean2[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.MeanFunction, NoneType, mfn.Constant, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_meanfunction__constant(
    p: Gaussian,
    mean1: mfn.MeanFunction,
    _: None,
    mean2: mfn.MeanFunction,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: General mean function
        - m2(.) :: Constant mean function

    :return: NxQ1xQ2
    """
    e_mean1 = expectation(p, mean1)
    return e_mean1[:, :, None] * mean2(p.mu)[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.Identity, NoneType, mfn.Identity, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, D, D]",
)
def _expectation_gaussian_identity__identity(
    p: Gaussian, mean1: mfn.Identity, _: None, mean2: mfn.Identity, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.), m2(.) :: Identity mean functions

    :return: NxDxD
    """
    return p.cov + (p.mu[:, :, None] * p.mu[:, None, :])


@dispatch.expectation.register(Gaussian, mfn.Identity, NoneType, mfn.Linear, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, D, Q]",
)
def _expectation_gaussian_identity__linear(
    p: Gaussian, mean1: mfn.Identity, _: None, mean2: mfn.Linear, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: Identity mean function
        - m2(.) :: Linear mean function

    :return: NxDxQ
    """
    N = tf.shape(p.mu)[0]
    e_xxt = p.cov + (p.mu[:, :, None] * p.mu[:, None, :])  # NxDxD
    e_xxt_A = tf.linalg.matmul(e_xxt, tf.tile(mean2.A[None, ...], (N, 1, 1)))  # NxDxQ
    e_x_bt = p.mu[:, :, None] * mean2.b[None, None, :]  # NxDxQ

    return e_xxt_A + e_x_bt


@dispatch.expectation.register(Gaussian, mfn.Linear, NoneType, mfn.Identity, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q, D]",
)
def _expectation_gaussian_linear__identity(
    p: Gaussian, mean1: mfn.Linear, _: None, mean2: mfn.Identity, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.) :: Linear mean function
        - m2(.) :: Identity mean function

    :return: NxQxD
    """
    N = tf.shape(p.mu)[0]
    e_xxt = p.cov + (p.mu[:, :, None] * p.mu[:, None, :])  # NxDxD
    e_A_xxt = tf.linalg.matmul(
        tf.tile(mean1.A[None, ...], (N, 1, 1)), e_xxt, transpose_a=True
    )  # NxQxD
    e_b_xt = mean1.b[None, :, None] * p.mu[:, None, :]  # NxQxD

    return e_A_xxt + e_b_xt


@dispatch.expectation.register(Gaussian, mfn.Linear, NoneType, mfn.Linear, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N, Q1, Q2]",
)
def _expectation_gaussian_linear__linear(
    p: Gaussian, mean1: mfn.Linear, _: None, mean2: mfn.Linear, __: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m1(x_n)^T m2(x_n)>_p(x_n)
        - m1(.), m2(.) :: Linear mean functions

    :return: NxQ1xQ2
    """
    e_xxt = p.cov + (p.mu[:, :, None] * p.mu[:, None, :])  # NxDxD
    e_A1t_xxt_A2 = tf.einsum("iq,nij,jz->nqz", mean1.A, e_xxt, mean2.A)  # NxQ1xQ2
    e_A1t_x_b2t = tf.einsum("iq,ni,z->nqz", mean1.A, p.mu, mean2.b)  # NxQ1xQ2
    e_b1_xt_A2 = tf.einsum("q,ni,iz->nqz", mean1.b, p.mu, mean2.A)  # NxQ1xQ2
    e_b1_b2t = mean1.b[:, None] * mean2.b[None, :]  # Q1xQ2

    return e_A1t_xxt_A2 + e_A1t_x_b2t + e_b1_xt_A2 + e_b1_b2t


# gpflow/expectations/misc.py

# 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 typing import Optional, Type, Union, cast

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints, InducingVariables
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from . import dispatch
from .expectations import ExpectationObject, PackedExpectationObject, expectation

NoneType: Type[None] = type(None)

# ================ exKxz transpose and mean function handling =================


@dispatch.expectation.register(
    (Gaussian, MarkovGaussian), mfn.Identity, NoneType, kernels.Linear, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_gaussian__linear_inducingpoints(
    p: Union[Gaussian, MarkovGaussian],
    mean: mfn.Identity,
    _: None,
    kernel: kernels.Linear,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <x_n K_{x_n, Z}>_p(x_n)
        - K_{.,} :: Linear kernel
    or the equivalent for MarkovGaussian

    :return: NxDxM
    """
    return tf.linalg.adjoint(expectation(p, (kernel, inducing_variable), mean))


@dispatch.expectation.register(
    (Gaussian, MarkovGaussian), kernels.Kernel, InducingVariables, mfn.MeanFunction, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M, Q]",
)
def _expectation_gaussian_kernel_inducingvariables__meanfunction(
    p: Union[Gaussian, MarkovGaussian],
    kernel: kernels.Kernel,
    inducing_variable: InducingVariables,
    mean: mfn.MeanFunction,
    _: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <K_{Z, x_n} m(x_n)>_p(x_n)
    or the equivalent for MarkovGaussian

    :return: NxMxQ
    """
    return tf.linalg.adjoint(expectation(p, mean, (kernel, inducing_variable), nghp=nghp))


@dispatch.expectation.register(Gaussian, mfn.Constant, NoneType, kernels.Kernel, InducingPoints)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, Q, M]",
)
def _expectation_gaussian_constant__kernel_inducingpoints(
    p: Gaussian,
    constant_mean: mfn.Constant,
    _: None,
    kernel: kernels.Kernel,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m(x_n)^T K_{x_n, Z}>_p(x_n)
        - m(x_i) = c :: Constant function
        - K_{.,.}    :: Kernel function

    :return: NxQxM
    """
    c = constant_mean(p.mu)  # NxQ
    eKxz = expectation(p, (kernel, inducing_variable), nghp=nghp)  # NxM

    return c[..., None] * eKxz[:, None, :]


@dispatch.expectation.register(Gaussian, mfn.Linear, NoneType, kernels.Kernel, InducingPoints)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, Q, M]",
)
def _expectation_gaussian_linear__kernel_inducingpoints(
    p: Gaussian,
    linear_mean: mfn.Linear,
    _: None,
    kernel: kernels.Kernel,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <m(x_n)^T K_{x_n, Z}>_p(x_n)
        - m(x_i) = A x_i + b :: Linear mean function
        - K_{.,.}            :: Kernel function

    :return: NxQxM
    """
    N = tf.shape(p.mu)[0]
    D = tf.shape(p.mu)[1]
    exKxz = expectation(p, mfn.Identity(D), (kernel, inducing_variable), nghp=nghp)
    eKxz = expectation(p, (kernel, inducing_variable), nghp=nghp)
    eAxKxz = tf.linalg.matmul(
        tf.tile(linear_mean.A[None, :, :], (N, 1, 1)), exKxz, transpose_a=True
    )
    ebKxz = linear_mean.b[None, :, None] * eKxz[:, None, :]
    return eAxKxz + ebKxz


@dispatch.expectation.register(Gaussian, mfn.Identity, NoneType, kernels.Kernel, InducingPoints)
def _expectation_gaussian__kernel_inducingpoints(
    p: Gaussian,
    identity_mean: mfn.Identity,
    _: None,
    kernel: kernels.Kernel,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    This prevents infinite recursion for kernels that don't have specific
    implementations of _expectation(p, identity_mean, None, kernel, inducing_variable).
    Recursion can arise because Identity is a subclass of Linear mean function
    so _expectation(p, linear_mean, none, kernel, inducing_variable) would call itself.
    More specific signatures (e.g. (p, identity_mean, None, RBF, inducing_variable)) will
    be found and used whenever available
    """
    raise NotImplementedError


# ============== Conversion to Gaussian from Diagonal or Markov ===============
# Catching missing DiagonalGaussian implementations by converting to full Gaussian:


@dispatch.expectation.register(
    DiagonalGaussian, object, (InducingVariables, NoneType), object, (InducingVariables, NoneType)
)
def _expectation_diagonal_generic(
    p: DiagonalGaussian,
    obj1: ExpectationObject,
    feat1: Optional[InducingVariables],
    obj2: ExpectationObject,
    feat2: Optional[InducingVariables],
    nghp: None = None,
) -> tf.Tensor:
    gaussian = Gaussian(p.mu, tf.linalg.diag(p.cov))
    return expectation(
        gaussian,
        cast(PackedExpectationObject, (obj1, feat1)),
        cast(PackedExpectationObject, (obj2, feat2)),
        nghp=nghp,
    )


# Catching missing MarkovGaussian implementations by converting to Gaussian (when indifferent):


@dispatch.expectation.register(
    MarkovGaussian, object, (InducingVariables, NoneType), object, (InducingVariables, NoneType)
)
def _expectation_markov_generic(
    p: MarkovGaussian,
    obj1: ExpectationObject,
    feat1: Optional[InducingVariables],
    obj2: ExpectationObject,
    feat2: Optional[InducingVariables],
    nghp: None = None,
) -> tf.Tensor:
    """
    Nota Bene: if only one object is passed, obj1 is
    associated with x_n, whereas obj2 with x_{n+1}

    """
    if obj2 is None:
        gaussian = Gaussian(p.mu[:-1], p.cov[0, :-1])
        return expectation(gaussian, cast(PackedExpectationObject, (obj1, feat1)), nghp=nghp)
    elif obj1 is None:
        gaussian = Gaussian(p.mu[1:], p.cov[0, 1:])
        return expectation(gaussian, cast(PackedExpectationObject, (obj2, feat2)), nghp=nghp)
    else:
        return expectation(
            p,
            cast(PackedExpectationObject, (obj1, feat1)),
            cast(PackedExpectationObject, (obj2, feat2)),
            nghp=nghp,
        )


# gpflow/expectations/products.py

# 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 functools import reduce
from typing import Type

import tensorflow as tf

from .. import kernels
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(DiagonalGaussian, kernels.Product, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_diagonal_product(
    p: DiagonalGaussian, kernel: kernels.Product, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\HadamardProd_i diag(Ki_{X[:, active_dims_i], X[:, active_dims_i]})>_p(X)
        - \HadamardProd_i Ki_{.,.} :: Product kernel
        - p                        :: DiagonalGaussian distribution (p.cov NxD)

    :return: N
    """
    if not kernel.on_separate_dimensions:
        raise NotImplementedError(
            "Product currently needs to be defined on separate dimensions."
        )  # pragma: no cover

    exps = [expectation(p, k, nghp=nghp) for k in kernel.kernels]
    return reduce(tf.multiply, exps)


@dispatch.expectation.register(
    DiagonalGaussian, kernels.Product, InducingPoints, NoneType, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_diagonal_product_inducingpoints(
    p: DiagonalGaussian,
    kernel: kernels.Product,
    inducing_variable: InducingPoints,
    __: None,
    ___: None,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\HadamardProd_i Ki_{X[:, active_dims_i], Z[:, active_dims_i]}>_p(X)
        - \HadamardProd_i Ki_{.,.} :: Product kernel
        - p                        :: DiagonalGaussian distribution (p.cov NxD)

    :return: NxM
    """
    if not kernel.on_separate_dimensions:
        raise NotImplementedError(
            "Product currently needs to be defined on separate dimensions."
        )  # pragma: no cover

    exps = [expectation(p, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.multiply, exps)


@dispatch.expectation.register(
    DiagonalGaussian, kernels.Product, InducingPoints, kernels.Product, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M, D, P]",
    "feat2: [M, D, P]",
    "return: [N, M, M]",
)
def _expectation_diagonal_product_inducingpoints__product_inducingpoints(
    p: DiagonalGaussian,
    kern1: kernels.Product,
    feat1: InducingPoints,
    kern2: kernels.Product,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = < prodK_{Z, x_n} prodK_{x_n, Z} >_p(x_n)
                   = < (\HadamardProd_i Ki_{Z[:, active_dims_i], x[n, active_dims_i]})  <-- Mx1
               1xM -->  (\HadamardProd_j Kj_{x[n, active_dims_j], Z[:, active_dims_j]}) >_p(x_n)  (MxM)

        - \HadamardProd_i Ki_{.,.}, \HadamardProd_j Kj_{.,.} :: Product kernels
        - p                        :: DiagonalGaussian distribution (p.cov NxD)

    :return: NxMxM
    """
    if feat1 != feat2:
        raise NotImplementedError("Different inducing variables are not supported.")
    if kern1 != kern2:
        raise NotImplementedError(
            "Calculating the expectation over two " "different Product kernels is not supported."
        )

    kernel = kern1
    inducing_variable = feat1

    if not kernel.on_separate_dimensions:
        raise NotImplementedError(
            "Product currently needs to be defined on separate dimensions."
        )  # pragma: no cover

    exps = [
        expectation(p, (k, inducing_variable), (k, inducing_variable), nghp=nghp)
        for k in kernel.kernels
    ]
    return reduce(tf.multiply, exps)


# gpflow/expectations/quadratures.py

# 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 typing import Any, Callable, Optional, Type, Union, cast

import numpy as np
import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..base import TensorType
from ..covariances import Kuf
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingVariables
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from ..quadrature import mvnquad
from . import dispatch
from .expectations import ExpectationObject, PackedExpectationObject, quadrature_expectation

register = dispatch.quadrature_expectation.register


NoneType: Type[None] = type(None)
EllipsisType = Any


def get_eval_func(
    obj: ExpectationObject,
    inducing_variable: Optional[InducingVariables],
    slice: Union[slice, EllipsisType, None] = None,
) -> Callable[[TensorType], tf.Tensor]:
    """
    Return the function of interest (kernel or mean) for the expectation
    depending on the type of :obj: and whether any inducing are given
    """

    slice = ... if slice is None else slice
    if inducing_variable is not None:
        # kernel + inducing_variable combination
        if not isinstance(inducing_variable, InducingVariables) or not isinstance(
            obj, kernels.Kernel
        ):
            raise TypeError("If `inducing_variable` is supplied, `obj` must be a kernel.")
        return lambda x: tf.transpose(Kuf(inducing_variable, obj, x))[slice]
    elif isinstance(obj, mfn.MeanFunction):
        return lambda x: obj(x)[slice]  # type: ignore[misc]
    elif isinstance(obj, kernels.Kernel):
        return lambda x: obj(x, full_cov=False)  # type: ignore[call-arg, misc]

    raise NotImplementedError()


@dispatch.quadrature_expectation.register(
    (Gaussian, DiagonalGaussian),
    object,
    (InducingVariables, NoneType),
    object,
    (InducingVariables, NoneType),
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable1: [M1, D, P]",
    "inducing_variable2: [M2, D, P]",
    "return: [N, ...]",
)
def _quadrature_expectation_gaussian(
    p: Union[Gaussian, DiagonalGaussian],
    obj1: ExpectationObject,
    inducing_variable1: Optional[InducingVariables],
    obj2: ExpectationObject,
    inducing_variable2: Optional[InducingVariables],
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    General handling of quadrature expectations for Gaussians and DiagonalGaussians
    Fallback method for missing analytic expectations
    """
    nghp = 100 if nghp is None else nghp

    # logger.warning(
    #     "Quadrature is used to calculate the expectation. This means that "
    #     "an analytical implementations is not available for the given combination."
    # )

    if obj1 is None:
        raise NotImplementedError("First object cannot be None.")

    if not isinstance(p, DiagonalGaussian):
        cov = p.cov
    else:
        if (
            isinstance(obj1, kernels.Kernel)
            and isinstance(obj2, kernels.Kernel)
            and obj1.on_separate_dims(obj2)
        ):  # no joint expectations required
            eKxz1 = quadrature_expectation(
                p, cast(PackedExpectationObject, (obj1, inducing_variable1)), nghp=nghp
            )
            eKxz2 = quadrature_expectation(
                p, cast(PackedExpectationObject, (obj2, inducing_variable2)), nghp=nghp
            )
            return eKxz1[:, :, None] * eKxz2[:, None, :]
        cov = tf.linalg.diag(p.cov)

    if obj2 is None:

        def eval_func(x: TensorType) -> tf.Tensor:
            fn = get_eval_func(obj1, inducing_variable1)
            return fn(x)

    else:

        def eval_func(x: TensorType) -> tf.Tensor:
            fn1 = get_eval_func(obj1, inducing_variable1, np.s_[:, :, None])
            fn2 = get_eval_func(obj2, inducing_variable2, np.s_[:, None, :])
            return fn1(x) * fn2(x)

    return mvnquad(eval_func, p.mu, cov, nghp)


@dispatch.quadrature_expectation.register(
    MarkovGaussian, object, (InducingVariables, NoneType), object, (InducingVariables, NoneType)
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable1: [M1, D, P]",
    "inducing_variable2: [M2, D, P]",
    "return: [N, ...]",
)
def _quadrature_expectation_markov(
    p: MarkovGaussian,
    obj1: ExpectationObject,
    inducing_variable1: Optional[InducingVariables],
    obj2: ExpectationObject,
    inducing_variable2: Optional[InducingVariables],
    nghp: Optional[int] = None,
) -> tf.Tensor:
    """
    Handling of quadrature expectations for Markov Gaussians (useful for time series)
    Fallback method for missing analytic expectations wrt Markov Gaussians
    Nota Bene: obj1 is always associated with x_n, whereas obj2 always with x_{n+1}
               if one requires e.g. <x_{n+1} K_{x_n, Z}>_p(x_{n:n+1}), compute the
               transpose and then transpose the result of the expectation
    """
    nghp = 40 if nghp is None else nghp

    # logger.warning(
    #     "Quadrature is used to calculate the expectation. This means that "
    #     "an analytical implementations is not available for the given combination."
    # )

    if obj2 is None:

        def eval_func(x: TensorType) -> tf.Tensor:
            return get_eval_func(obj1, inducing_variable1)(x)

        mu, cov = p.mu[:-1], p.cov[0, :-1]  # cross covariances are not needed
    elif obj1 is None:

        def eval_func(x: TensorType) -> tf.Tensor:
            return get_eval_func(obj2, inducing_variable2)(x)

        mu, cov = p.mu[1:], p.cov[0, 1:]  # cross covariances are not needed
    else:

        def eval_func(x: TensorType) -> tf.Tensor:
            x1 = tf.split(x, 2, 1)[0]
            x2 = tf.split(x, 2, 1)[1]
            res1 = get_eval_func(obj1, inducing_variable1, np.s_[:, :, None])(x1)
            res2 = get_eval_func(obj2, inducing_variable2, np.s_[:, None, :])(x2)
            return res1 * res2

        mu = tf.concat((p.mu[:-1, :], p.mu[1:, :]), 1)  # Nx2D
        cov_top = tf.concat((p.cov[0, :-1, :, :], p.cov[1, :-1, :, :]), 2)  # NxDx2D
        cov_bottom = tf.concat((tf.linalg.adjoint(p.cov[1, :-1, :, :]), p.cov[0, 1:, :, :]), 2)
        cov = tf.concat((cov_top, cov_bottom), 1)  # Nx2Dx2D

    return mvnquad(eval_func, mu, cov, nghp)


# gpflow/expectations/squared_exponentials.py

# 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 typing import Type, Union

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from ..utilities.ops import square_distance
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, kernels.SquaredExponential, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_gaussian_sqe(
    p: Gaussian, kernel: kernels.SquaredExponential, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    """
    Compute the expectation:
    <diag(K_{X, X})>_p(X)
        - K_{.,.} :: RBF kernel

    :return: N
    """
    return kernel(p.mu, full_cov=False)


@dispatch.expectation.register(
    Gaussian, kernels.SquaredExponential, InducingPoints, NoneType, NoneType
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_gaussian_sqe_inducingpoints(
    p: Gaussian,
    kernel: kernels.SquaredExponential,
    inducing_variable: InducingPoints,
    _: None,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    <K_{X, Z}>_p(X)
        - K_{.,.} :: RBF kernel

    :return: NxM
    """
    # use only active dimensions
    Xcov = kernel.slice_cov(p.cov)
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)
    D = tf.shape(Xmu)[1]

    lengthscales = kernel.lengthscales
    if not kernel.ard:
        lengthscales = tf.zeros((D,), dtype=lengthscales.dtype) + kernel.lengthscales

    chol_L_plus_Xcov = tf.linalg.cholesky(tf.linalg.diag(lengthscales ** 2) + Xcov)  # NxDxD

    all_diffs = tf.transpose(Z) - tf.expand_dims(Xmu, 2)  # NxDxM
    exponent_mahalanobis = tf.linalg.triangular_solve(
        chol_L_plus_Xcov, all_diffs, lower=True
    )  # NxDxM
    exponent_mahalanobis = tf.reduce_sum(tf.square(exponent_mahalanobis), 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    sqrt_det_L = tf.reduce_prod(lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N

    return kernel.variance * (determinants[:, None] * exponent_mahalanobis)


@dispatch.expectation.register(
    Gaussian, mfn.Identity, NoneType, kernels.SquaredExponential, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_gaussian__sqe_inducingpoints(
    p: Gaussian,
    mean: mfn.Identity,
    _: None,
    kernel: kernels.SquaredExponential,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <x_n K_{x_n, Z}>_p(x_n)
        - K_{.,.} :: RBF kernel

    :return: NxDxM
    """
    Xmu, Xcov = p.mu, p.cov

    D = tf.shape(Xmu)[1]

    lengthscales = kernel.lengthscales
    if not kernel.ard:
        lengthscales = tf.zeros((D,), dtype=lengthscales.dtype) + lengthscales

    chol_L_plus_Xcov = tf.linalg.cholesky(tf.linalg.diag(lengthscales ** 2) + Xcov)  # NxDxD
    all_diffs = tf.transpose(inducing_variable.Z) - tf.expand_dims(Xmu, 2)  # NxDxM

    sqrt_det_L = tf.reduce_prod(lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N

    exponent_mahalanobis = tf.linalg.cholesky_solve(chol_L_plus_Xcov, all_diffs)  # NxDxM
    non_exponent_term = tf.linalg.matmul(Xcov, exponent_mahalanobis, transpose_a=True)
    non_exponent_term = tf.expand_dims(Xmu, 2) + non_exponent_term  # NxDxM

    exponent_mahalanobis = tf.reduce_sum(all_diffs * exponent_mahalanobis, 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    return (
        kernel.variance
        * (determinants[:, None] * exponent_mahalanobis)[:, None, :]
        * non_exponent_term
    )


@dispatch.expectation.register(
    MarkovGaussian, mfn.Identity, NoneType, kernels.SquaredExponential, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_markov__sqe_inducingpoints(
    p: MarkovGaussian,
    mean: mfn.Identity,
    _: None,
    kernel: kernels.SquaredExponential,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <x_{n+1} K_{x_n, Z}>_p(x_{n:n+1})
        - K_{.,.} :: RBF kernel
        - p       :: MarkovGaussian distribution (p.cov 2x(N+1)xDxD)

    :return: NxDxM
    """
    Xmu, Xcov = p.mu, p.cov

    D = tf.shape(Xmu)[1]
    lengthscales = kernel.lengthscales
    if not kernel.ard:
        lengthscales = tf.zeros((D,), dtype=lengthscales.dtype) + lengthscales

    chol_L_plus_Xcov = tf.linalg.cholesky(tf.linalg.diag(lengthscales ** 2) + Xcov[0, :-1])  # NxDxD
    all_diffs = tf.transpose(inducing_variable.Z) - tf.expand_dims(Xmu[:-1], 2)  # NxDxM

    sqrt_det_L = tf.reduce_prod(lengthscales)
    sqrt_det_L_plus_Xcov = tf.exp(
        tf.reduce_sum(tf.math.log(tf.linalg.diag_part(chol_L_plus_Xcov)), axis=1)
    )
    determinants = sqrt_det_L / sqrt_det_L_plus_Xcov  # N

    exponent_mahalanobis = tf.linalg.cholesky_solve(chol_L_plus_Xcov, all_diffs)  # NxDxM
    non_exponent_term = tf.linalg.matmul(Xcov[1, :-1], exponent_mahalanobis, transpose_a=True)
    non_exponent_term = tf.expand_dims(Xmu[1:], 2) + non_exponent_term  # NxDxM

    exponent_mahalanobis = tf.reduce_sum(all_diffs * exponent_mahalanobis, 1)  # NxM
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxM

    return (
        kernel.variance
        * (determinants[:, None] * exponent_mahalanobis)[:, None, :]
        * non_exponent_term
    )


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian),
    kernels.SquaredExponential,
    InducingPoints,
    kernels.SquaredExponential,
    InducingPoints,
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M, D, P]",
    "feat2: [M, D, P]",
    "return: [N, M, M]",
)
def _expectation_gaussian_sqe_inducingpoints__sqe_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    kern1: kernels.SquaredExponential,
    feat1: InducingPoints,
    kern2: kernels.SquaredExponential,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    """
    Compute the expectation:
    expectation[n] = <Ka_{Z1, x_n} Kb_{x_n, Z2}>_p(x_n)
        - Ka_{.,.}, Kb_{.,.} :: RBF kernels
    Ka and Kb as well as Z1 and Z2 can differ from each other, but this is supported
    only if the Gaussian p is Diagonal (p.cov NxD) and Ka, Kb have disjoint active_dims
    in which case the joint expectations simplify into a product of expectations

    :return: NxMxM
    """
    if kern1.on_separate_dims(kern2) and isinstance(
        p, DiagonalGaussian
    ):  # no joint expectations required
        eKxz1 = expectation(p, (kern1, feat1))
        eKxz2 = expectation(p, (kern2, feat2))
        return eKxz1[:, :, None] * eKxz2[:, None, :]

    if feat1 != feat2 or kern1 != kern2:
        raise NotImplementedError(
            "The expectation over two kernels has only an "
            "analytical implementation if both kernels are equal."
        )

    kernel = kern1
    inducing_variable = feat1

    # use only active dimensions
    Xcov = kernel.slice_cov(tf.linalg.diag(p.cov) if isinstance(p, DiagonalGaussian) else p.cov)
    Z, Xmu = kernel.slice(inducing_variable.Z, p.mu)

    N = tf.shape(Xmu)[0]
    D = tf.shape(Xmu)[1]

    squared_lengthscales = kernel.lengthscales ** 2
    if not kernel.ard:
        zero_lengthscales = tf.zeros((D,), dtype=squared_lengthscales.dtype)
        squared_lengthscales = squared_lengthscales + zero_lengthscales

    sqrt_det_L = tf.reduce_prod(0.5 * squared_lengthscales) ** 0.5
    C = tf.linalg.cholesky(0.5 * tf.linalg.diag(squared_lengthscales) + Xcov)  # NxDxD
    dets = sqrt_det_L / tf.exp(tf.reduce_sum(tf.math.log(tf.linalg.diag_part(C)), axis=1))  # N

    C_inv_mu = tf.linalg.triangular_solve(C, tf.expand_dims(Xmu, 2), lower=True)  # NxDx1
    C_inv_z = tf.linalg.triangular_solve(
        C, tf.tile(tf.expand_dims(0.5 * tf.transpose(Z), 0), [N, 1, 1]), lower=True
    )  # NxDxM
    mu_CC_inv_mu = tf.expand_dims(tf.reduce_sum(tf.square(C_inv_mu), 1), 2)  # Nx1x1
    z_CC_inv_z = tf.reduce_sum(tf.square(C_inv_z), 1)  # NxM
    zm_CC_inv_zn = tf.linalg.matmul(C_inv_z, C_inv_z, transpose_a=True)  # NxMxM
    two_z_CC_inv_mu = 2 * tf.linalg.matmul(C_inv_z, C_inv_mu, transpose_a=True)[:, :, 0]  # NxM
    # NxMxM
    exponent_mahalanobis = (
        mu_CC_inv_mu
        + tf.expand_dims(z_CC_inv_z, 1)
        + tf.expand_dims(z_CC_inv_z, 2)
        + 2 * zm_CC_inv_zn
        - tf.expand_dims(two_z_CC_inv_mu, 2)
        - tf.expand_dims(two_z_CC_inv_mu, 1)
    )
    exponent_mahalanobis = tf.exp(-0.5 * exponent_mahalanobis)  # NxMxM

    # Compute sqrt(self(Z)) explicitly to prevent automatic gradient from
    # being NaN sometimes, see pull request #615
    kernel_sqrt = tf.exp(-0.25 * square_distance(Z / kernel.lengthscales, None))
    return kernel.variance ** 2 * kernel_sqrt * tf.reshape(dets, [N, 1, 1]) * exponent_mahalanobis


# gpflow/expectations/sums.py

# 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.

import itertools
from functools import reduce
from typing import Type, Union

import tensorflow as tf

from .. import kernels
from .. import mean_functions as mfn
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints
from ..probability_distributions import DiagonalGaussian, Gaussian, MarkovGaussian
from . import dispatch
from .expectations import expectation

NoneType: Type[None] = type(None)


@dispatch.expectation.register(Gaussian, kernels.Sum, NoneType, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "return: [N]",
)
def _expectation_gaussian_sum(
    p: Gaussian, kernel: kernels.Sum, _: None, __: None, ___: None, nghp: None = None
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\Sum_i diag(Ki_{X, X})>_p(X)
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: N
    """
    exps = [expectation(p, k, nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(Gaussian, kernels.Sum, InducingPoints, NoneType, NoneType)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, M]",
)
def _expectation_gaussian_sum_inducingpoints(
    p: Gaussian,
    kernel: kernels.Sum,
    inducing_variable: InducingPoints,
    _: None,
    __: None,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    <\Sum_i Ki_{X, Z}>_p(X)
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: NxM
    """
    exps = [expectation(p, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(
    Gaussian, (mfn.Linear, mfn.Identity, mfn.Constant), NoneType, kernels.Sum, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, Q, M]",
)
def _expectation_gaussian_linear__sum_inducingpoints(
    p: Gaussian,
    mean: Union[mfn.Linear, mfn.Identity, mfn.Constant],
    _: None,
    kernel: kernels.Sum,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = <m(x_n)^T (\Sum_i Ki_{x_n, Z})>_p(x_n)
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: NxQxM
    """
    exps = [expectation(p, mean, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(MarkovGaussian, mfn.Identity, NoneType, kernels.Sum, InducingPoints)
@check_shapes(
    "p: [N, D]",
    "inducing_variable: [M, D, P]",
    "return: [N, D, M]",
)
def _expectation_markov__sum_inducingpoints(
    p: MarkovGaussian,
    mean: mfn.Identity,
    _: None,
    kernel: kernels.Sum,
    inducing_variable: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = <x_{n+1} (\Sum_i Ki_{x_n, Z})>_p(x_{n:n+1})
        - \Sum_i Ki_{.,.} :: Sum kernel

    :return: NxDxM
    """
    exps = [expectation(p, mean, (k, inducing_variable), nghp=nghp) for k in kernel.kernels]
    return reduce(tf.add, exps)


@dispatch.expectation.register(
    (Gaussian, DiagonalGaussian), kernels.Sum, InducingPoints, kernels.Sum, InducingPoints
)
@check_shapes(
    "p: [N, D]",
    "feat1: [M1, D, P]",
    "feat2: [M2, D, P]",
    "return: [N, M1, M2]",
)
def _expectation_gaussian_sum_inducingpoints__sum_inducingpoints(
    p: Union[Gaussian, DiagonalGaussian],
    kern1: kernels.Sum,
    feat1: InducingPoints,
    kern2: kernels.Sum,
    feat2: InducingPoints,
    nghp: None = None,
) -> tf.Tensor:
    r"""
    Compute the expectation:
    expectation[n] = <(\Sum_i K1_i_{Z1, x_n}) (\Sum_j K2_j_{x_n, Z2})>_p(x_n)
        - \Sum_i K1_i_{.,.}, \Sum_j K2_j_{.,.} :: Sum kernels

    :return: NxM1xM2
    """
    crossexps = []

    if kern1 == kern2 and feat1 == feat2:  # avoid duplicate computation by using transposes
        for i, k1 in enumerate(kern1.kernels):
            crossexps.append(expectation(p, (k1, feat1), (k1, feat1), nghp=nghp))

            for k2 in kern1.kernels[:i]:
                eKK = expectation(p, (k1, feat1), (k2, feat2), nghp=nghp)
                eKK += tf.linalg.adjoint(eKK)
                crossexps.append(eKK)
    else:
        for k1, k2 in itertools.product(kern1.kernels, kern2.kernels):
            crossexps.append(expectation(p, (k1, feat1), (k2, feat2), nghp=nghp))

    return reduce(tf.add, crossexps)


# gpflow/experimental/__init__.py

# Copyright 2016-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.
"""
This package contains code that is still experimental. This means:

* Expect breaking changes.
* Expect poor documentation.
* Expect bugs.

When the code is considered mature it will be moved out of this package.
"""
from . import check_shapes, utils

__all__ = [
    "check_shapes",
    "utils",
]


# gpflow/experimental/utils.py

# Copyright 2016-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 functools import wraps
from typing import Any, Callable, TypeVar, cast
from warnings import warn

C = TypeVar("C", bound=Callable[..., Any])


def experimental(func: C) -> C:
    """
    Decorator that marks the decorated function as experimental.

    The first time an experimental function is called, a warning is printed.

    Example::

        @experimental
        def forty_two() -> int:
            return 42
    """

    has_warned = False

    @wraps(func)
    def wrap_experimental(*args: Any, **kwargs: Any) -> Any:
        nonlocal has_warned

        if not has_warned:
            name = f"{func.__module__}.{func.__qualname__}"
            warn(
                f"You're calling {name} which is considered *experimental*."
                " Expect: breaking changes, poor documentation, and bugs."
            )
            has_warned = True

        return func(*args, **kwargs)

    return cast(C, wrap_experimental)


# gpflow/inducing_variables/__init__.py

from . import multioutput
from .inducing_patch import InducingPatches
from .inducing_variables import InducingPoints, InducingVariables, Multiscale
from .multioutput import (
    FallbackSeparateIndependentInducingVariables,
    FallbackSharedIndependentInducingVariables,
    MultioutputInducingVariables,
    SeparateIndependentInducingVariables,
    SharedIndependentInducingVariables,
)

__all__ = [
    "FallbackSeparateIndependentInducingVariables",
    "FallbackSharedIndependentInducingVariables",
    "InducingPatches",
    "InducingPoints",
    "InducingVariables",
    "MultioutputInducingVariables",
    "Multiscale",
    "SeparateIndependentInducingVariables",
    "SharedIndependentInducingVariables",
    "inducing_patch",
    "inducing_variables",
    "multioutput",
]


# gpflow/inducing_variables/inducing_patch.py

# 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 .inducing_variables import InducingPoints


class InducingPatches(InducingPoints):
    pass


# gpflow/inducing_variables/inducing_variables.py

# 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.

import abc
from typing import Optional

import tensorflow as tf
import tensorflow_probability as tfp
from deprecated import deprecated

from ..base import Module, Parameter, TensorData, TensorType
from ..experimental.check_shapes import ErrorContext, Shape, check_shapes, register_get_shape
from ..utilities import positive


class InducingVariables(Module, abc.ABC):
    """
    Abstract base class for inducing variables.
    """

    @property
    @abc.abstractmethod
    def num_inducing(self) -> tf.Tensor:
        """
        Returns the number of inducing variables, relevant for example to determine the size of the
        variational distribution.
        """
        raise NotImplementedError

    @deprecated(
        reason="len(iv) should return an `int`, but this actually returns a `tf.Tensor`."
        " Use `iv.num_inducing` instead."
    )
    def __len__(self) -> tf.Tensor:
        return self.num_inducing

    @property
    @abc.abstractmethod
    def shape(self) -> Shape:
        """
        Return the shape of these inducing variables.

        Shape should be some variation of ``[M, D, P]``, where:

        * ``M`` is the number of inducing variables.
        * ``D`` is the number of input dimensions.
        * ``P`` is the number of output dimensions (1 if this is not a multi-output inducing
          variable).
        """


@register_get_shape(InducingVariables)
def get_scalar_shape(shaped: InducingVariables, context: ErrorContext) -> Shape:
    return shaped.shape


class InducingPointsBase(InducingVariables):
    @check_shapes(
        "Z: [M, D]",
    )
    def __init__(self, Z: TensorData, name: Optional[str] = None):
        """
        :param Z: The initial positions of the inducing points.
        """
        super().__init__(name=name)
        if not isinstance(Z, (tf.Variable, tfp.util.TransformedVariable)):
            Z = Parameter(Z)
        self.Z = Z

    @property  # type: ignore[misc]  # mypy doesn't like decorated properties.
    @check_shapes(
        "return: []",
    )
    def num_inducing(self) -> Optional[tf.Tensor]:
        return tf.shape(self.Z)[0]

    @property
    def shape(self) -> Shape:
        shape = self.Z.shape
        if not shape:
            return None
        return tuple(shape) + (1,)


class InducingPoints(InducingPointsBase):
    """
    Real-space inducing points
    """


class Multiscale(InducingPointsBase):
    r"""
    Multi-scale inducing variables

    Originally proposed in :cite:t:`NIPS2009_3876`.
    """

    @check_shapes(
        "Z: [M, D]",
        "scales: [M, D]",
    )
    def __init__(self, Z: TensorData, scales: TensorData):
        super().__init__(Z)
        # Multi-scale inducing_variable widths (std. dev. of Gaussian)
        self.scales = Parameter(scales, transform=positive())

    @staticmethod
    @check_shapes(
        "A: [N, D]",
        "B: [M, D]",
        "sc: [broadcast N, broadcast M, D]",
        "return: [N, M]",
    )
    def _cust_square_dist(A: TensorType, B: TensorType, sc: TensorType) -> tf.Tensor:
        """
        Custom version of _square_dist that allows sc to provide per-datapoint length
        scales.
        """
        return tf.reduce_sum(tf.square((tf.expand_dims(A, 1) - tf.expand_dims(B, 0)) / sc), 2)


# gpflow/kernels/__init__.py

from . import multioutput
from .base import Combination, Kernel, Product, Sum
from .changepoints import ChangePoints
from .convolutional import Convolutional
from .linears import Linear, Polynomial
from .misc import ArcCosine, Coregion
from .multioutput import (
    IndependentLatent,
    LinearCoregionalization,
    MultioutputKernel,
    SeparateIndependent,
    SharedIndependent,
)
from .periodic import Periodic
from .statics import Constant, Static, White

from .stationaries import (  # isort:skip
    # base classes:
    Stationary,
    IsotropicStationary,
    AnisotropicStationary,
    # actual kernel classes:
    Cosine,
    Exponential,
    Matern12,
    Matern32,
    Matern52,
    RationalQuadratic,
    SquaredExponential,
)

Bias = Constant
RBF = SquaredExponential

__all__ = [
    "AnisotropicStationary",
    "ArcCosine",
    "Bias",
    "ChangePoints",
    "Combination",
    "Constant",
    "Convolutional",
    "Coregion",
    "Cosine",
    "Exponential",
    "IndependentLatent",
    "IsotropicStationary",
    "Kernel",
    "Linear",
    "LinearCoregionalization",
    "Matern12",
    "Matern32",
    "Matern52",
    "MultioutputKernel",
    "Periodic",
    "Polynomial",
    "Product",
    "RBF",
    "RationalQuadratic",
    "SeparateIndependent",
    "SharedIndependent",
    "SquaredExponential",
    "Static",
    "Stationary",
    "Sum",
    "White",
    "base",
    "changepoints",
    "convolutional",
    "linears",
    "misc",
    "multioutput",
    "periodic",
    "statics",
    "stationaries",
]


# gpflow/kernels/base.py

# Copyright 2018-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.

r"""
Kernels form a core component of GPflow models and allow prior information to
be encoded about a latent function of interest. The effect of choosing
different kernels, and how it is possible to combine multiple kernels is shown
in the `"Using kernels in GPflow" notebook <notebooks/kernels.html>`_.

Broadcasting over leading dimensions:
`kernel.K(X1, X2)` returns the kernel evaluated on every pair in X1 and X2.
E.g. if X1 has shape [S1, N1, D] and X2 has shape [S2, N2, D], kernel.K(X1, X2)
will return a tensor of shape [S1, N1, S2, N2]. Similarly, kernel.K(X1, X1)
returns a tensor of shape [S1, N1, S1, N1]. In contrast, the return shape of
kernel.K(X1) is [S1, N1, N1]. (Without leading dimensions, the behaviour of
kernel.K(X, None) is identical to kernel.K(X, X).)
"""

import abc
from functools import partial, reduce
from typing import Callable, List, Optional, Sequence, Tuple, Union, overload

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray, Module, TensorType
from ..experimental.check_shapes import check_shapes

ActiveDims = Union[slice, Sequence[int]]
NormalizedActiveDims = Union[slice, AnyNDArray]


class Kernel(Module, metaclass=abc.ABCMeta):
    """
    The basic kernel class. Handles active dims.
    """

    def __init__(
        self, active_dims: Optional[ActiveDims] = None, name: Optional[str] = None
    ) -> None:
        """
        :param active_dims: active dimensions, either a slice or list of
            indices into the columns of X.
        :param name: optional kernel name.
        """
        super().__init__(name=name)
        self._active_dims = self._normalize_active_dims(active_dims)

    @staticmethod
    def _normalize_active_dims(value: Optional[ActiveDims]) -> NormalizedActiveDims:
        if value is None:
            return slice(None, None, None)
        if isinstance(value, slice):
            return value
        return np.array(value, dtype=int)

    @property
    def active_dims(self) -> NormalizedActiveDims:
        return self._active_dims

    @active_dims.setter
    def active_dims(self, value: ActiveDims) -> None:
        self._active_dims = self._normalize_active_dims(value)

    def on_separate_dims(self, other: "Kernel") -> bool:
        """
        Checks if the dimensions, over which the kernels are specified, overlap.
        Returns True if they are defined on different/separate dimensions and False otherwise.
        """
        if isinstance(self.active_dims, slice) or isinstance(other.active_dims, slice):
            # Be very conservative for kernels defined over slices of dimensions
            return False

        if self.active_dims is None or other.active_dims is None:
            return False

        this_dims = self.active_dims.reshape(-1, 1)
        other_dims = other.active_dims.reshape(1, -1)
        return not np.any(this_dims == other_dims)

    @overload
    def slice(self, X: TensorType, X2: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        ...

    @overload
    def slice(self, X: TensorType, X2: None) -> Tuple[tf.Tensor, None]:
        ...

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return[0]: [batch..., N, I]",
        "return[1]: [batch2..., N2, I]",
    )
    def slice(
        self, X: TensorType, X2: Optional[TensorType] = None
    ) -> Tuple[tf.Tensor, Optional[tf.Tensor]]:
        """
        Slice the correct dimensions for use in the kernel, as indicated by `self.active_dims`.

        :param X: Input 1.
        :param X2: Input 2, can be None.
        :return: Sliced X, X2.
        """
        dims = self.active_dims
        if isinstance(dims, slice):
            X = X[..., dims]
            if X2 is not None:
                X2 = X2[..., dims]
        elif dims is not None:
            X = tf.gather(X, dims, axis=-1)
            if X2 is not None:
                X2 = tf.gather(X2, dims, axis=-1)
        return X, X2

    @check_shapes(
        "cov: [N, *D_or_DD]",
        "return: [N, I, I]",
    )
    def slice_cov(self, cov: TensorType) -> tf.Tensor:
        """
        Slice the correct dimensions for use in the kernel, as indicated by
        `self.active_dims` for covariance matrices. This requires slicing the
        rows *and* columns. This will also turn flattened diagonal
        matrices into a tensor of full diagonal matrices.

        :param cov: Tensor of covariance matrices.
        :return: Sliced covariance matrices.
        """
        cov = tf.convert_to_tensor(cov)
        assert isinstance(cov, tf.Tensor)  # Hint for mypy.

        if cov.shape.ndims == 2:
            cov = tf.linalg.diag(cov)

        dims = self.active_dims

        if isinstance(dims, slice):
            return cov[..., dims, dims]
        elif dims is not None:
            nlast = tf.shape(cov)[-1]
            ndims = len(dims)

            cov_shape = tf.shape(cov)
            cov_reshaped = tf.reshape(cov, [-1, nlast, nlast])
            gather1 = tf.gather(tf.transpose(cov_reshaped, [2, 1, 0]), dims)
            gather2 = tf.gather(tf.transpose(gather1, [1, 0, 2]), dims)
            cov = tf.reshape(
                tf.transpose(gather2, [2, 0, 1]), tf.concat([cov_shape[:-2], [ndims, ndims]], 0)
            )

        return cov

    @check_shapes(
        "ard_parameter: [any...]",
    )
    def _validate_ard_active_dims(self, ard_parameter: TensorType) -> None:
        """
        Validate that ARD parameter matches the number of active_dims (provided active_dims
        has been specified as an array).
        """
        ard_parameter = tf.convert_to_tensor(ard_parameter)
        assert isinstance(ard_parameter, tf.Tensor)  # Hint for mypy.

        if self.active_dims is None or isinstance(self.active_dims, slice):
            # Can only validate parameter if active_dims is an array
            return

        if ard_parameter.shape.rank > 0 and ard_parameter.shape[0] != len(self.active_dims):
            raise ValueError(
                f"Size of `active_dims` {self.active_dims} does not match "
                f"size of ard parameter ({ard_parameter.shape[0]})"
            )

    @abc.abstractmethod
    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if X2 is not None",
        "return: [batch..., N, N] if X2 is None",
    )
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        raise NotImplementedError

    @abc.abstractmethod
    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N]",
    )
    def K_diag(self, X: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if full_cov and (X2 is not None)",
        "return: [batch..., N, N] if full_cov and (X2 is None)",
        "return: [batch..., N] if not full_cov",
    )
    def __call__(
        self,
        X: TensorType,
        X2: Optional[TensorType] = None,
        *,
        full_cov: bool = True,
        presliced: bool = False,
    ) -> tf.Tensor:
        if (not full_cov) and (X2 is not None):
            raise ValueError("Ambiguous inputs: `not full_cov` and `X2` are not compatible.")

        if not presliced:
            X, X2 = self.slice(X, X2)

        if not full_cov:
            assert X2 is None
            return self.K_diag(X)

        else:
            return self.K(X, X2)

    def __add__(self, other: "Kernel") -> "Kernel":
        return Sum([self, other])

    def __mul__(self, other: "Kernel") -> "Kernel":
        return Product([self, other])


class Combination(Kernel):
    """
    Combine a list of kernels, e.g. by adding or multiplying (see inheriting
    classes).

    The names of the kernels to be combined are generated from their class
    names.
    """

    _reduction = None

    def __init__(self, kernels: Sequence[Kernel], name: Optional[str] = None) -> None:
        super().__init__(name=name)

        if not all(isinstance(k, Kernel) for k in kernels):
            raise TypeError("can only combine Kernel instances")  # pragma: no cover

        self.kernels: List[Kernel] = []
        self._set_kernels(kernels)

    def _set_kernels(self, kernels: Sequence[Kernel]) -> None:
        # add kernels to a list, flattening out instances of this class therein
        kernels_list: List[Kernel] = []
        for k in kernels:
            if isinstance(k, self.__class__):
                kernels_list.extend(k.kernels)
            else:
                kernels_list.append(k)
        self.kernels = kernels_list

    @property
    def on_separate_dimensions(self) -> bool:
        """
        Checks whether the kernels in the combination act on disjoint subsets
        of dimensions. Currently, it is hard to asses whether two slice objects
        will overlap, so this will always return False.

        :return: Boolean indicator.
        """
        if np.any([isinstance(k.active_dims, slice) for k in self.kernels]):
            # Be conservative in the case of a slice object
            return False
        else:
            dimlist = [k.active_dims for k in self.kernels]
            overlapping = False
            for i, dims_i in enumerate(dimlist):
                assert isinstance(dims_i, np.ndarray)
                for dims_j in dimlist[i + 1 :]:
                    assert isinstance(dims_j, np.ndarray)
                    if np.any(dims_i.reshape(-1, 1) == dims_j.reshape(1, -1)):
                        overlapping = True
            return not overlapping


class ReducingCombination(Combination):
    def __call__(
        self,
        X: TensorType,
        X2: Optional[TensorType] = None,
        *,
        full_cov: bool = True,
        presliced: bool = False,
    ) -> tf.Tensor:
        return self._reduce(
            [k(X, X2, full_cov=full_cov, presliced=presliced) for k in self.kernels]
        )

    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        return self._reduce([k.K(X, X2) for k in self.kernels])

    def K_diag(self, X: TensorType) -> tf.Tensor:
        return self._reduce([k.K_diag(X) for k in self.kernels])

    @property
    @abc.abstractmethod
    def _reduce(self) -> Callable[[Sequence[TensorType]], TensorType]:
        pass


class Sum(ReducingCombination):
    @property
    def _reduce(self) -> Callable[[Sequence[TensorType]], TensorType]:
        return tf.add_n  # type: ignore[no-any-return]


class Product(ReducingCombination):
    @property
    def _reduce(self) -> Callable[[Sequence[TensorType]], TensorType]:
        return partial(reduce, tf.multiply)


# gpflow/kernels/changepoints.py

# 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 typing import Optional, Sequence

import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shape as cs
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import Combination, Kernel


class ChangePoints(Combination):
    r"""
    The ChangePoints kernel defines a fixed number of change-points along a 1d
    input space where different kernels govern different parts of the space.

    The kernel is by multiplication and addition of the base kernels with
    sigmoid functions (σ). A single change-point kernel is defined as::

        K₁(x, x') * (1 - σ(x)) * (1 - σ(x')) + K₂(x, x') * σ(x) * σ(x')

    where K₁ is deactivated around the change-point and K₂ is activated. The
    single change-point version can be found in :cite:t:`lloyd2014`. Each sigmoid
    is a logistic function defined as::

        σ(x) = 1 / (1 + exp{-s(x - x₀)})

    parameterized by location "x₀" and steepness "s".

    The key reference is :cite:t:`lloyd2014`.
    """

    @check_shapes(
        "locations: [n_change_points]",
        "steepness: [broadcast n_change_points]",
    )
    def __init__(
        self,
        kernels: Sequence[Kernel],
        locations: TensorType,
        steepness: TensorType = 1.0,
        name: Optional[str] = None,
    ):
        """
        :param kernels: list of kernels defining the different regimes
        :param locations: list of change-point locations in the 1d input space
        :param steepness: the steepness parameter(s) of the sigmoids, this can be
            common between them or decoupled
        """
        if len(kernels) != len(locations) + 1:
            raise ValueError(
                "Number of kernels ({nk}) must be one more than the number of "
                "changepoint locations ({nl})".format(nk=len(kernels), nl=len(locations))
            )

        if isinstance(steepness, Sequence) and len(steepness) != len(locations):
            raise ValueError(
                "Dimension of steepness ({ns}) does not match number of changepoint "
                "locations ({nl})".format(ns=len(steepness), nl=len(locations))
            )

        super().__init__(kernels, name=name)

        self.locations = Parameter(locations)
        self.steepness = Parameter(steepness, transform=positive())

    def _set_kernels(self, kernels: Sequence[Kernel]) -> None:
        # it is not clear how to flatten out nested change-points
        self.kernels = list(kernels)

    @inherit_check_shapes
    def K(self, X: tf.Tensor, X2: Optional[tf.Tensor] = None) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `ChangePoints` kernel requires a 1D input space.")

        rank = tf.rank(X) - 2
        batch = tf.shape(X)[:-2]
        N = tf.shape(X)[-2]
        Ncp = tf.shape(self.locations)[0]
        sig_X = cs(self._sigmoids(X), "[batch..., N, 1, Ncp]")

        if X2 is None:
            rank2 = 0
            batch2 = tf.constant([], dtype=tf.int32)
            N2 = N
            sig_X2 = sig_X
            sig_X = cs(
                tf.reshape(sig_X, tf.concat([batch, [N, 1, Ncp]], 0)), "[batch..., N, 1, Ncp]"
            )
            sig_X2 = cs(
                tf.reshape(sig_X2, tf.concat([batch, [1, N, Ncp]], 0)), "[batch..., 1, N, Ncp]"
            )
        else:
            rank2 = tf.rank(X2) - 2
            batch2 = tf.shape(X2)[:-2]
            N2 = tf.shape(X2)[-2]

            sig_X2 = cs(self._sigmoids(X2), "[batch2..., N2, 1, Ncp]")

            ones = tf.ones((rank,), dtype=tf.int32)
            ones2 = tf.ones((rank2,), dtype=tf.int32)
            sig_X = cs(
                tf.reshape(sig_X, tf.concat([batch, [N], ones2, [1, Ncp]], 0)),
                "[batch..., N, ..., 1, Ncp]",
            )
            sig_X2 = cs(
                tf.reshape(sig_X2, tf.concat([ones, [1], batch2, [N2, Ncp]], 0)),
                "[..., 1, batch2..., N2, Ncp]",
            )

        # `starters` are the sigmoids going from 0 -> 1, whilst `stoppers` go
        # from 1 -> 0.
        starters = cs(sig_X * sig_X2, "[batch..., N, batch2..., N2, Ncp]")
        stoppers = cs((1 - sig_X) * (1 - sig_X2), "[batch..., N, batch2..., N2, Ncp]")

        # prepend `starters` with ones and append ones to `stoppers` since the
        # first kernel has no start and the last kernel has no end
        ones = tf.ones(tf.concat([batch, [N], batch2, [N2, 1]], 0), dtype=X.dtype)
        starters = cs(tf.concat([ones, starters], axis=-1), "[batch..., N, batch2..., N2, Nkern]")
        stoppers = cs(tf.concat([stoppers, ones], axis=-1), "[batch..., N, batch2..., N2, Nkern]")

        # now combine with the underlying kernels
        kernel_stack = cs(
            tf.stack([k(X, X2) for k in self.kernels], axis=-1),
            "[batch..., N, batch2..., N2, Nkern]",
        )
        return tf.reduce_sum(kernel_stack * starters * stoppers, axis=-1)

    @inherit_check_shapes
    def K_diag(self, X: tf.Tensor) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `ChangePoints` kernel requires a 1D input space.")

        batch = tf.shape(X)[:-2]
        N = tf.shape(X)[-2]
        Ncp = tf.shape(self.locations)[0]

        sig_X = cs(
            tf.reshape(self._sigmoids(X), tf.concat([batch, [N, Ncp]], 0)), "[batch..., N, Ncp]"
        )

        ones = tf.ones(tf.concat([batch, [N, 1]], 0), dtype=X.dtype)
        starters = cs(tf.concat([ones, sig_X * sig_X], axis=-1), "[batch..., N, Nkern]")
        stoppers = cs(tf.concat([(1 - sig_X) * (1 - sig_X), ones], axis=-1), "[batch..., N, Nkern]")

        kernel_stack = cs(
            tf.stack([k(X, full_cov=False) for k in self.kernels], axis=-1), "[batch..., N, Nkern]"
        )
        return tf.reduce_sum(kernel_stack * starters * stoppers, axis=-1)

    @check_shapes(
        "X: [batch...]",
        "return: [batch..., Ncp]",
    )
    def _sigmoids(self, X: tf.Tensor) -> tf.Tensor:
        locations = tf.sort(self.locations)  # ensure locations are ordered
        locations = tf.reshape(locations, (-1,))
        steepness = tf.reshape(self.steepness, (-1,))
        return tf.sigmoid(steepness * (X[..., None] - locations))


# gpflow/kernels/convolutional.py

# 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 typing import Optional, Sequence, cast

import numpy as np
import tensorflow as tf

from ..base import Parameter, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shape as cs
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import to_default_float
from .base import Kernel


class Convolutional(Kernel):
    r"""
    Plain convolutional kernel as described in :cite:t:`vdw2017convgp`. Defines
    a GP :math:`f()` that is constructed from a sum of responses of individual patches
    in an image:

    .. math::
       f(x) = \sum_p x^{[p]}

    where :math:`x^{[p]}` is the :math:`p`'th patch in the image.

    The key reference is :cite:t:`vdw2017convgp`.
    """

    @check_shapes(
        "image_shape: [2]",
        "patch_shape: [2]",
        "weights: [P]",
    )
    def __init__(
        self,
        base_kernel: Kernel,
        image_shape: Sequence[int],
        patch_shape: Sequence[int],
        weights: Optional[TensorType] = None,
        colour_channels: int = 1,
    ) -> None:
        super().__init__()
        self.image_shape = image_shape
        self.patch_shape = patch_shape
        self.base_kernel = base_kernel
        self.colour_channels = colour_channels
        self.weights = Parameter(
            np.ones(self.num_patches, dtype=default_float()) if weights is None else weights
        )

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N, P, S]",
    )
    def get_patches(self, X: TensorType) -> tf.Tensor:
        """
        Extracts patches from the images X. Patches are extracted separately for each of the colour
        channels.

        :param X: Images.
        :return: Patches.
        """
        # Roll the colour channel to the front, so it appears to
        # `tf.extract_image_patches()` as separate images. Then extract patches
        # and reshape to have the first axis the same as the number of images.
        # The separate patches will then be in the second axis.
        batch = tf.shape(X)[:-2]
        N = tf.shape(X)[-2]
        flat_batch = tf.reduce_prod(batch)
        num_data = flat_batch * N
        X = cs(
            tf.transpose(tf.reshape(X, [num_data, -1, self.colour_channels]), [0, 2, 1]),
            "[num_data, C, W_x_H]",
        )
        X = cs(
            tf.reshape(X, [-1, self.image_shape[0], self.image_shape[1], 1], name="rX"),
            "[num_data_x_C, W, H, 1]",
        )
        patches = cs(
            tf.image.extract_patches(
                X,
                [1, self.patch_shape[0], self.patch_shape[1], 1],
                [1, 1, 1, 1],
                [1, 1, 1, 1],
                "VALID",
            ),
            "[num_data_x_C, n_x_patches, n_y_patches, S]",
        )
        shp = tf.shape(patches)
        reshaped_patches = cs(
            tf.reshape(
                patches, tf.concat([batch, [N, self.colour_channels * shp[1] * shp[2], shp[3]]], 0)
            ),
            "[batch..., N, P, S]",
        )
        return to_default_float(reshaped_patches)

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        Xp = cs(self.get_patches(X), "[batch..., N, P, S]")
        W2 = cs(self.weights[:, None] * self.weights[None, :], "[P, P]")

        rank = tf.rank(Xp) - 3
        batch = tf.shape(Xp)[:-3]
        N = tf.shape(Xp)[-3]
        P = tf.shape(Xp)[-2]
        S = tf.shape(Xp)[-1]
        ones = tf.ones((rank,), dtype=tf.int32)

        if X2 is None:
            Xp = cs(tf.reshape(Xp, tf.concat([batch, [N * P, S]], 0)), "[batch..., N_x_P, S]")
            bigK = cs(self.base_kernel.K(Xp), "[batch..., N_x_P, N_x_P]")
            bigK = cs(
                tf.reshape(bigK, tf.concat([batch, [N, P, N, P]], 0)), "[batch..., N, P, N, P]"
            )
            W2 = cs(tf.reshape(W2, tf.concat([ones, [1, P, 1, P]], 0)), "[..., 1, P, 1, P]")
            W2bigK = cs(bigK * W2, "[batch..., N, P, N, P]")
            return cs(
                tf.reduce_sum(W2bigK, [rank + 1, rank + 3]) / self.num_patches ** 2.0,
                "[batch..., N, N]",
            )

        else:
            Xp2 = Xp if X2 is None else cs(self.get_patches(X2), "[batch2..., N2, P, S]")
            rank2 = tf.rank(Xp2) - 3
            ones2 = tf.ones((rank2,), dtype=tf.int32)
            bigK = cs(self.base_kernel.K(Xp, Xp2), "[batch..., N, P, batch2..., N2, P]")
            W2 = cs(
                tf.reshape(W2, tf.concat([ones, [1, P], ones2, [1, P]], 0)),
                "[..., 1, P, ..., 1, P]",
            )
            W2bigK = cs(bigK * W2, "[batch..., N, P, batch2..., N2, P]")
            return cs(
                tf.reduce_sum(W2bigK, [rank + 1, rank + rank2 + 3]) / self.num_patches ** 2.0,
                "[batch..., N, batch2..., N2]",
            )

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        Xp = cs(self.get_patches(X), "[batch..., N, P, S]")

        rank = tf.rank(Xp) - 3
        P = tf.shape(Xp)[-2]
        ones = tf.ones((rank,), dtype=tf.int32)

        W2 = cs(self.weights[:, None] * self.weights[None, :], "[P, P]")
        W2 = cs(tf.reshape(W2, tf.concat([ones, [1, P, P]], 0)), "[..., 1, P, P]")

        bigK = cs(self.base_kernel.K(Xp), "[batch..., N, P, P]")

        return tf.reduce_sum(bigK * W2, [rank + 1, rank + 2]) / self.num_patches ** 2.0

    @property
    def patch_len(self) -> int:
        return cast(int, np.prod(self.patch_shape))

    @property
    def num_patches(self) -> int:
        return (
            (self.image_shape[0] - self.patch_shape[0] + 1)
            * (self.image_shape[1] - self.patch_shape[1] + 1)
            * self.colour_channels
        )


# gpflow/kernels/linears.py

# 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 typing import Optional

import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import ActiveDims, Kernel


class Linear(Kernel):
    """
    The linear kernel. Functions drawn from a GP with this kernel are linear, i.e. f(x) = cx.
    The kernel equation is

        k(x, y) = σ²xy

    where σ² is the variance parameter.
    """

    @check_shapes(
        "variance: [broadcast n_active_dims]",
    )
    def __init__(
        self, variance: TensorType = 1.0, active_dims: Optional[ActiveDims] = None
    ) -> None:
        """
        :param variance: the (initial) value for the variance parameter(s),
            to induce ARD behaviour this must be initialised as an array the same
            length as the the number of active dimensions e.g. [1., 1., 1.]
        :param active_dims: a slice or list specifying which columns of X are used
        """
        super().__init__(active_dims)
        self.variance = Parameter(variance, transform=positive())
        self._validate_ard_active_dims(self.variance)

    @property
    def ard(self) -> bool:
        """
        Whether ARD behaviour is active.
        """
        ndims: int = self.variance.shape.ndims
        return ndims > 0

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        if X2 is None:
            return tf.matmul(X * self.variance, X, transpose_b=True)
        else:
            return tf.tensordot(X * self.variance, X2, [[-1], [-1]])

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return tf.reduce_sum(tf.square(X) * self.variance, axis=-1)


class Polynomial(Linear):
    """
    The Polynomial kernel. Functions drawn from a GP with this kernel are
    polynomials of degree `d`. The kernel equation is

        k(x, y) = (σ²xy + γ)ᵈ

    where:
    σ² is the variance parameter,
    γ is the offset parameter,
    d is the degree parameter.
    """

    @check_shapes(
        "variance: [broadcast n_active_dims]",
    )
    def __init__(
        self,
        degree: TensorType = 3.0,
        variance: TensorType = 1.0,
        offset: TensorType = 1.0,
        active_dims: Optional[ActiveDims] = None,
    ) -> None:
        """
        :param degree: the degree of the polynomial
        :param variance: the (initial) value for the variance parameter(s),
            to induce ARD behaviour this must be initialised as an array the same
            length as the the number of active dimensions e.g. [1., 1., 1.]
        :param offset: the offset of the polynomial
        :param active_dims: a slice or list specifying which columns of X are used
        """
        super().__init__(variance, active_dims)
        self.degree = degree
        self.offset = Parameter(offset, transform=positive())

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        return (super().K(X, X2) + self.offset) ** self.degree

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return (super().K_diag(X) + self.offset) ** self.degree


# gpflow/kernels/misc.py

# Copyright 2018-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 typing import Optional

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray, Parameter, TensorType
from ..experimental.check_shapes import check_shape as cs
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive, to_default_float
from .base import ActiveDims, Kernel


class ArcCosine(Kernel):
    """
    The Arc-cosine family of kernels which mimics the computation in neural
    networks. The order parameter specifies the assumed activation function.
    The Multi Layer Perceptron (MLP) kernel is closely related to the ArcCosine
    kernel of order 0.

    The key reference is :cite:t:`NIPS2009_3628`.
    """

    implemented_orders = {0, 1, 2}

    @check_shapes(
        "variance: []",
        "weight_variances: [broadcast n_active_dims]",
        "bias_variance: []",
    )
    def __init__(
        self,
        order: int = 0,
        variance: TensorType = 1.0,
        weight_variances: TensorType = 1.0,
        bias_variance: TensorType = 1.0,
        *,
        active_dims: Optional[ActiveDims] = None,
        name: Optional[str] = None,
    ) -> None:
        """
        :param order: specifies the activation function of the neural network
          the function is a rectified monomial of the chosen order
        :param variance: the (initial) value for the variance parameter
        :param weight_variances: the (initial) value for the weight_variances parameter,
            to induce ARD behaviour this must be initialised as an array the same
            length as the the number of active dimensions e.g. [1., 1., 1.]
        :param bias_variance: the (initial) value for the bias_variance parameter
            defaults to 1.0
        :param active_dims: a slice or list specifying which columns of X are used
        """
        super().__init__(active_dims=active_dims, name=name)

        if order not in self.implemented_orders:
            raise ValueError("Requested kernel order is not implemented.")
        self.order = order

        self.variance = Parameter(variance, transform=positive())
        self.bias_variance = Parameter(bias_variance, transform=positive())
        self.weight_variances = Parameter(weight_variances, transform=positive())
        self._validate_ard_active_dims(self.weight_variances)

    @property
    def ard(self) -> bool:
        """
        Whether ARD behaviour is active.
        """
        ndims: int = self.weight_variances.shape.ndims
        return ndims > 0

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N]",
    )
    def _diag_weighted_product(self, X: TensorType) -> tf.Tensor:
        return tf.reduce_sum(self.weight_variances * tf.square(X), axis=-1) + self.bias_variance

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if X2 is not None",
        "return: [batch..., N, N] if X2 is None",
    )
    def _full_weighted_product(self, X: TensorType, X2: Optional[TensorType]) -> tf.Tensor:
        if X2 is None:
            return (
                tf.linalg.matmul((self.weight_variances * X), X, transpose_b=True)
                + self.bias_variance
            )

        else:
            D = tf.shape(X)[-1]
            N = tf.shape(X)[-2]
            N2 = tf.shape(X2)[-2]
            batch = tf.shape(X)[:-2]
            batch2 = tf.shape(X2)[:-2]
            rank = tf.rank(X) - 2
            rank2 = tf.rank(X2) - 2
            ones = tf.ones((rank,), tf.int32)
            ones2 = tf.ones((rank2,), tf.int32)

            X = cs(
                tf.reshape(X, tf.concat([batch, ones2, [N, D]], 0)),
                "[batch..., broadcast batch2..., N, D]",
            )
            X2 = cs(
                tf.reshape(X2, tf.concat([ones, batch2, [N2, D]], 0)),
                "[broadcast batch..., batch2..., N2, D]",
            )
            result = cs(
                tf.linalg.matmul((self.weight_variances * X), X2, transpose_b=True)
                + self.bias_variance,
                "[batch..., batch2..., N, N2]",
            )

            indices = tf.concat(
                [
                    tf.range(rank),
                    [rank + rank2],
                    tf.range(rank2) + rank,
                    [rank + rank2 + 1],
                ],
                axis=0,
            )
            return tf.transpose(result, indices)

    @check_shapes(
        "theta: [any...]",
        "return: [any...]",
    )
    def _J(self, theta: TensorType) -> TensorType:
        """
        Implements the order dependent family of functions defined in equations
        4 to 7 in the reference paper.
        """
        if self.order == 0:
            return np.pi - theta
        elif self.order == 1:
            return tf.sin(theta) + (np.pi - theta) * tf.cos(theta)
        else:
            assert self.order == 2, f"Don't know how to handle order {self.order}."
            return 3.0 * tf.sin(theta) * tf.cos(theta) + (np.pi - theta) * (
                1.0 + 2.0 * tf.cos(theta) ** 2
            )

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        X_denominator = cs(tf.sqrt(self._diag_weighted_product(X)), "[batch..., N]")
        if X2 is None:
            X2_denominator = cs(X_denominator[..., None, :], "[batch..., 1, N]")
            X_denominator = cs(X_denominator[..., :, None], "[batch..., N, 1]")
            numerator = cs(self._full_weighted_product(X, None), "[batch..., N, N]")
        else:
            X2_denominator = cs(tf.sqrt(self._diag_weighted_product(X2)), "[batch2..., N2]")

            batch = tf.shape(X)[:-1]
            batch2 = tf.shape(X2)[:-1]
            ones = tf.ones((tf.rank(X) - 1,), tf.int32)
            ones2 = tf.ones((tf.rank(X2) - 1,), tf.int32)

            X_denominator = cs(
                tf.reshape(X_denominator, tf.concat([batch, ones2], 0)),
                "[batch..., N, broadcast batch2..., 1]",
            )
            X2_denominator = cs(
                tf.reshape(X2_denominator, tf.concat([ones, batch2], 0)),
                "[broadcast batch..., 1, batch2..., N2]",
            )
            numerator = cs(self._full_weighted_product(X, X2), "[batch..., N, batch2..., N2]")

        cos_theta = numerator / X_denominator / X2_denominator
        jitter = 1e-15
        theta = tf.acos(jitter + (1 - 2 * jitter) * cos_theta)

        return (
            self.variance
            * (1.0 / np.pi)
            * self._J(theta)
            * X_denominator ** self.order
            * X2_denominator ** self.order
        )

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        X_product = self._diag_weighted_product(X)
        const = (1.0 / np.pi) * self._J(to_default_float(0.0))
        return self.variance * const * X_product ** self.order


class Coregion(Kernel):
    """
    A Coregionalization kernel. The inputs to this kernel are _integers_ (we
    cast them from floats as needed) which usually specify the *outputs* of a
    Coregionalization model.

    The kernel function is an indexing of a positive-definite matrix:

      K(x, y) = B[x, y] .

    To ensure that B is positive-definite, it is specified by the two
    parameters of this kernel, W and kappa:

      B = W Wᵀ + diag(kappa) .

    We refer to the size of B as "output_dim x output_dim", since this is the
    number of outputs in a coregionalization model. We refer to the number of
    columns on W as 'rank': it is the number of degrees of correlation between
    the outputs.

    NB. There is a symmetry between the elements of W, which creates a local
    minimum at W=0. To avoid this, it is recommended to initialize the
    optimization (or MCMC chain) using a random W.
    """

    def __init__(
        self,
        output_dim: int,
        rank: int,
        *,
        active_dims: Optional[ActiveDims] = None,
        name: Optional[str] = None,
    ) -> None:
        """
        :param output_dim: number of outputs expected (0 <= X < output_dim)
        :param rank: number of degrees of correlation between outputs
        """

        # assert input_dim == 1, "Coregion kernel in 1D only"
        super().__init__(active_dims=active_dims, name=name)

        self.output_dim = output_dim
        self.rank = rank
        W: AnyNDArray = 0.1 * np.ones((self.output_dim, self.rank))
        kappa = np.ones(self.output_dim)
        self.W = Parameter(W)
        self.kappa = Parameter(kappa, transform=positive())

    @check_shapes(
        "return: [P, P]",
    )
    def output_covariance(self) -> tf.Tensor:
        B = tf.linalg.matmul(self.W, self.W, transpose_b=True) + tf.linalg.diag(self.kappa)
        return B

    @check_shapes(
        "return: [P]",
    )
    def output_variance(self) -> tf.Tensor:
        B_diag = tf.reduce_sum(tf.square(self.W), 1) + self.kappa
        return B_diag

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        B = cs(self.output_covariance(), "[O, O]")
        X = cs(tf.cast(X[..., 0], tf.int32), "[batch..., N]")
        if X2 is None:
            batch = tf.shape(X)[:-1]
            N = tf.shape(X)[-1]
            O = tf.shape(B)[-1]

            result = cs(tf.gather(B, X), "[batch..., N, O]")
            result = cs(tf.reshape(result, [-1, N, O]), "[flat_batch, N, O]")
            flat_X = cs(tf.reshape(X, [-1, N]), "[flat_batch, N]")
            result = cs(tf.gather(result, flat_X, axis=2, batch_dims=1), "[flat_batch, N, N]")
            result = cs(tf.reshape(result, tf.concat([batch, [N, N]], 0)), "[batch..., N, N]")
        else:
            X2 = cs(tf.cast(X2[..., 0], tf.int32), "[batch2..., N2]")

            rank2 = tf.rank(X2)

            result = cs(tf.gather(B, X2), "[batch2..., N2, O]")
            result = cs(
                tf.transpose(result, tf.concat([[rank2], tf.range(rank2)], 0)), "[O, batch2..., N2]"
            )
            result = cs(tf.gather(result, X), "[batch..., N, batch2..., N2]")

        return result

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        cs(X, "[batch..., N, 1]  # The `Coregion` kernel requires a 1D input space.")

        X = tf.cast(X[..., 0], tf.int32)
        B_diag = self.output_variance()
        return tf.gather(B_diag, X)


# gpflow/kernels/periodic.py

# 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 typing import Optional

import numpy as np
import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from ..utilities.ops import difference_matrix
from .base import ActiveDims, Kernel, NormalizedActiveDims
from .stationaries import IsotropicStationary


class Periodic(Kernel):
    """
    The periodic family of kernels. Can be used to wrap any Stationary kernel
    to transform it into a periodic version. The canonical form (based on the
    SquaredExponential kernel) can be found in Equation (47) of

    D.J.C.MacKay. Introduction to Gaussian processes. In C.M.Bishop, editor,
    Neural Networks and Machine Learning, pages 133--165. Springer, 1998.

    The derivation can be achieved by mapping the original inputs through the
    transformation u = (cos(x), sin(x)).

    For the SquaredExponential base kernel, the result can be expressed as:

        k(r) =  σ² exp{ -0.5 sin²(π r / γ) / ℓ²}

    where:
    r is the Euclidean distance between the input points
    ℓ is the lengthscales parameter,
    σ² is the variance parameter,
    γ is the period parameter.

    NOTE: usually we have a factor of 4 instead of 0.5 in front but this
        is absorbed into the lengthscales hyperparameter.
    NOTE: periodic kernel uses `active_dims` of a base kernel, therefore
        the constructor doesn't have it as an argument.
    """

    @check_shapes(
        "period: [broadcast n_active_dims]",
    )
    def __init__(self, base_kernel: IsotropicStationary, period: TensorType = 1.0) -> None:
        """
        :param base_kernel: the base kernel to make periodic; must inherit from Stationary
            Note that `active_dims` should be specified in the base kernel.
        :param period: the period; to induce a different period per active dimension
            this must be initialized with an array the same length as the number
            of active dimensions e.g. [1., 1., 1.]
        """
        if not isinstance(base_kernel, IsotropicStationary):
            raise TypeError("Periodic requires an IsotropicStationary kernel as the `base_kernel`")

        super().__init__()
        self.base_kernel = base_kernel
        self.period = Parameter(period, transform=positive())
        self.base_kernel._validate_ard_active_dims(self.period)

    @property
    def active_dims(self) -> NormalizedActiveDims:
        return self.base_kernel.active_dims

    @active_dims.setter
    def active_dims(self, value: ActiveDims) -> None:
        # type-ignore below is because mypy doesn't understand that getter and the setter of
        # `active_dims` have different types.
        self.base_kernel.active_dims = value  # type: ignore[assignment]

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return self.base_kernel.K_diag(X)

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        r = np.pi * (difference_matrix(X, X2)) / self.period
        scaled_sine = tf.sin(r) / self.base_kernel.lengthscales
        if hasattr(self.base_kernel, "K_r"):
            sine_r = tf.reduce_sum(tf.abs(scaled_sine), -1)
            K = self.base_kernel.K_r(sine_r)
        else:
            sine_r2 = tf.reduce_sum(tf.square(scaled_sine), -1)
            K = self.base_kernel.K_r2(sine_r2)
        return K


# gpflow/kernels/statics.py

# 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 typing import Optional

import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import ActiveDims, Kernel


class Static(Kernel):
    """
    Kernels who don't depend on the value of the inputs are 'Static'.  The only
    parameter is a variance, σ².
    """

    @check_shapes(
        "variance: []",
    )
    def __init__(
        self, variance: TensorType = 1.0, active_dims: Optional[ActiveDims] = None
    ) -> None:
        super().__init__(active_dims)
        self.variance = Parameter(variance, transform=positive())

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return tf.fill(tf.shape(X)[:-1], tf.squeeze(self.variance))


class White(Static):
    """
    The White kernel: this kernel produces 'white noise'. The kernel equation is

        k(x_n, x_m) = δ(n, m) σ²

    where:
    δ(.,.) is the Kronecker delta,
    σ²  is the variance parameter.
    """

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        if X2 is None:
            d = tf.fill(tf.shape(X)[:-1], tf.squeeze(self.variance))
            return tf.linalg.diag(d)
        else:
            shape = tf.concat([tf.shape(X)[:-1], tf.shape(X2)[:-1]], axis=0)
            return tf.zeros(shape, dtype=X.dtype)


class Constant(Static):
    """
    The Constant (aka Bias) kernel. Functions drawn from a GP with this kernel
    are constant, i.e. f(x) = c, with c ~ N(0, σ^2). The kernel equation is

        k(x, y) = σ²

    where:
    σ²  is the variance parameter.
    """

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        if X2 is None:
            shape = tf.concat(
                [
                    tf.shape(X)[:-2],
                    tf.reshape(tf.shape(X)[-2], [1]),
                    tf.reshape(tf.shape(X)[-2], [1]),
                ],
                axis=0,
            )
        else:
            shape = tf.concat([tf.shape(X)[:-1], tf.shape(X2)[:-1]], axis=0)

        return tf.fill(shape, tf.squeeze(self.variance))


# gpflow/kernels/stationaries.py

# 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 typing import Any, Optional

import numpy as np
import tensorflow as tf

from ..base import Parameter, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from ..utilities.ops import difference_matrix, square_distance
from .base import ActiveDims, Kernel


class Stationary(Kernel):
    """
    Base class for kernels that are stationary, that is, they only depend on

        d = x - x'

    This class handles 'ard' behaviour, which stands for 'Automatic Relevance
    Determination'. This means that the kernel has one lengthscale per
    dimension, otherwise the kernel is isotropic (has a single lengthscale).
    """

    @check_shapes(
        "variance: []",
        "lengthscales: [broadcast n_active_dims]",
    )
    def __init__(
        self, variance: TensorType = 1.0, lengthscales: TensorType = 1.0, **kwargs: Any
    ) -> None:
        """
        :param variance: the (initial) value for the variance parameter.
        :param lengthscales: the (initial) value for the lengthscale
            parameter(s), to induce ARD behaviour this must be initialised as
            an array the same length as the number of active dimensions
            e.g. [1., 1., 1.]. If only a single value is passed, this value
            is used as the lengthscale of each dimension.
        :param kwargs: accepts `name` and `active_dims`, which is a list or
            slice of indices which controls which columns of X are used (by
            default, all columns are used).
        """
        for kwarg in kwargs:
            if kwarg not in {"name", "active_dims"}:
                raise TypeError(f"Unknown keyword argument: {kwarg}")

        super().__init__(**kwargs)
        self.variance = Parameter(variance, transform=positive())
        self.lengthscales = Parameter(lengthscales, transform=positive())
        self._validate_ard_active_dims(self.lengthscales)

    @property
    def ard(self) -> bool:
        """
        Whether ARD behaviour is active.
        """
        ndims: int = self.lengthscales.shape.ndims
        return ndims > 0

    @check_shapes(
        "X: [broadcast any...]",
        "return: [any...]",
    )
    def scale(self, X: TensorType) -> TensorType:
        X_scaled = X / self.lengthscales if X is not None else X
        return X_scaled

    @inherit_check_shapes
    def K_diag(self, X: TensorType) -> tf.Tensor:
        return tf.fill(tf.shape(X)[:-1], tf.squeeze(self.variance))


class IsotropicStationary(Stationary):
    """
    Base class for isotropic stationary kernels, i.e. kernels that only
    depend on

        r = ‖x - x'‖

    Derived classes should implement one of:

        K_r2(self, r2): Returns the kernel evaluated on r² (r2), which is the
        squared scaled Euclidean distance Should operate element-wise on r2.

        K_r(self, r): Returns the kernel evaluated on r, which is the scaled
        Euclidean distance. Should operate element-wise on r.
    """

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        r2 = self.scaled_squared_euclid_dist(X, X2)
        return self.K_r2(r2)

    @check_shapes(
        "r2: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r2(self, r2: TensorType) -> tf.Tensor:
        if hasattr(self, "K_r"):
            # Clipping around the (single) float precision which is ~1e-45.
            r = tf.sqrt(tf.maximum(r2, 1e-36))
            return self.K_r(r)  # pylint: disable=no-member
        raise NotImplementedError

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2] if X2 is not None",
        "return: [batch..., N, N] if X2 is None",
    )
    def scaled_squared_euclid_dist(
        self, X: TensorType, X2: Optional[TensorType] = None
    ) -> tf.Tensor:
        """
        Returns ‖(X - X2ᵀ) / ℓ‖², i.e. the squared L₂-norm.
        """
        return square_distance(self.scale(X), self.scale(X2))


class AnisotropicStationary(Stationary):
    """
    Base class for anisotropic stationary kernels, i.e. kernels that only
    depend on

        d = x - x'

    Derived classes should implement K_d(self, d): Returns the kernel evaluated
    on d, which is the pairwise difference matrix, scaled by the lengthscale
    parameter ℓ (i.e. [(X - X2ᵀ) / ℓ]). The last axis corresponds to the
    input dimension.
    """

    @check_shapes(
        "variance: []",
        "lengthscales: [broadcast n_active_dims]",
    )
    def __init__(
        self, variance: TensorType = 1.0, lengthscales: TensorType = 1.0, **kwargs: Any
    ) -> None:
        """
        :param variance: the (initial) value for the variance parameter.
        :param lengthscales: the (initial) value for the lengthscale
            parameter(s), to induce ARD behaviour this must be initialised as
            an array the same length as the number of active dimensions
            e.g. [1., 1., 1.]. Note that anisotropic kernels can possess
            negative lengthscales. If only a single value is passed, this
            value is used as the lengthscale of each dimension.
        :param kwargs: accepts `name` and `active_dims`, which is a list or
            slice of indices which controls which columns of X are used (by
            default, all columns are used).
        """
        super().__init__(variance, lengthscales, **kwargs)

        if self.ard:
            self.lengthscales = Parameter(self.lengthscales.numpy())

    @inherit_check_shapes
    def K(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        return self.K_d(self.scaled_difference_matrix(X, X2))

    @check_shapes(
        "X: [batch..., N, D]",
        "X2: [batch2..., N2, D]",
        "return: [batch..., N, batch2..., N2, D] if X2 is not None",
        "return: [batch..., N, N, D] if X2 is None",
    )
    def scaled_difference_matrix(self, X: TensorType, X2: Optional[TensorType] = None) -> tf.Tensor:
        """
        Returns [(X - X2ᵀ) / ℓ]. If X has shape [..., N, D] and
        X2 has shape [..., M, D], the output will have shape [..., N, M, D].
        """
        return difference_matrix(self.scale(X), self.scale(X2))

    @check_shapes(
        "d: [batch..., N, D]",
        "return: [batch..., N]",
    )
    def K_d(self, d: TensorType) -> tf.Tensor:
        raise NotImplementedError


class SquaredExponential(IsotropicStationary):
    """
    The radial basis function (RBF) or squared exponential kernel. The kernel equation is

        k(r) = σ² exp{-½ r²}

    where:
    r   is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ.
    σ²  is the variance parameter

    Functions drawn from a GP with this kernel are infinitely differentiable!
    """

    @inherit_check_shapes
    def K_r2(self, r2: TensorType) -> tf.Tensor:
        return self.variance * tf.exp(-0.5 * r2)


class RationalQuadratic(IsotropicStationary):
    """
    Rational Quadratic kernel,

    k(r) = σ² (1 + r² / 2αℓ²)^(-α)

    σ² : variance
    ℓ  : lengthscales
    α  : alpha, determines relative weighting of small-scale and large-scale fluctuations

    For α → ∞, the RQ kernel becomes equivalent to the squared exponential.
    """

    def __init__(
        self,
        variance: TensorType = 1.0,
        lengthscales: TensorType = 1.0,
        alpha: TensorType = 1.0,
        active_dims: Optional[ActiveDims] = None,
    ) -> None:
        super().__init__(variance=variance, lengthscales=lengthscales, active_dims=active_dims)
        self.alpha = Parameter(alpha, transform=positive())

    @inherit_check_shapes
    def K_r2(self, r2: TensorType) -> tf.Tensor:
        return self.variance * (1 + 0.5 * r2 / self.alpha) ** (-self.alpha)


class Exponential(IsotropicStationary):
    """
    The Exponential kernel. It is equivalent to a Matern12 kernel with doubled lengthscales.
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        return self.variance * tf.exp(-0.5 * r)


class Matern12(IsotropicStationary):
    """
    The Matern 1/2 kernel. Functions drawn from a GP with this kernel are not
    differentiable anywhere. The kernel equation is

    k(r) = σ² exp{-r}

    where:
    r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ.
    σ² is the variance parameter
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        return self.variance * tf.exp(-r)


class Matern32(IsotropicStationary):
    """
    The Matern 3/2 kernel. Functions drawn from a GP with this kernel are once
    differentiable. The kernel equation is

    k(r) = σ² (1 + √3r) exp{-√3 r}

    where:
    r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ,
    σ² is the variance parameter.
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        sqrt3 = np.sqrt(3.0)
        return self.variance * (1.0 + sqrt3 * r) * tf.exp(-sqrt3 * r)


class Matern52(IsotropicStationary):
    """
    The Matern 5/2 kernel. Functions drawn from a GP with this kernel are twice
    differentiable. The kernel equation is

    k(r) = σ² (1 + √5r + 5/3r²) exp{-√5 r}

    where:
    r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ,
    σ² is the variance parameter.
    """

    @check_shapes(
        "r: [batch..., N]",
        "return: [batch..., N]",
    )
    def K_r(self, r: TensorType) -> tf.Tensor:
        sqrt5 = np.sqrt(5.0)
        return self.variance * (1.0 + sqrt5 * r + 5.0 / 3.0 * tf.square(r)) * tf.exp(-sqrt5 * r)


class Cosine(AnisotropicStationary):
    """
    The Cosine kernel. Functions drawn from a GP with this kernel are sinusoids
    (with a random phase).  The kernel equation is

        k(r) = σ² cos{2πd}

    where:
    d  is the sum of the per-dimension differences between the input points, scaled by the
    lengthscale parameter ℓ (i.e. Σᵢ [(X - X2ᵀ) / ℓ]ᵢ),
    σ² is the variance parameter.
    """

    @inherit_check_shapes
    def K_d(self, d: TensorType) -> tf.Tensor:
        d = tf.reduce_sum(d, axis=-1)
        return self.variance * tf.cos(2 * np.pi * d)


# gpflow/likelihoods/__init__.py

# 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.
"""
Likelihoods are another core component of GPflow. This describes how likely the
data is under the assumptions made about the underlying latent functions
p(Y|F). Different likelihoods make different
assumptions about the distribution of the data, as such different data-types
(continuous, binary, ordinal, count) are better modelled with different
likelihood assumptions.

Use of any likelihood other than Gaussian typically introduces the need to use
an approximation to perform inference, if one isn't already needed.
Variational inference and MCMC models are included in GPflow and allow
approximate inference with non-Gaussian likelihoods. An introduction to these
models can be found :ref:`here <implemented_models>`. Specific notebooks
illustrating non-Gaussian likelihood regressions are available for
`classification <notebooks/classification.html>`_ (binary data), `ordinal
<notebooks/ordinal.html>`_ and `multiclass <notebooks/multiclass.html>`_.

Creating new likelihoods
------------------------

Likelihoods are defined by their
log-likelihood. When creating new likelihoods, the
:func:`logp <gpflow.likelihoods.Likelihood.logp>` method (log p(Y|F)), the
:func:`conditional_mean <gpflow.likelihoods.Likelihood.conditional_mean>`,
:func:`conditional_variance
<gpflow.likelihoods.Likelihood.conditional_variance>`.

In order to perform variational inference with non-Gaussian likelihoods a term
called ``variational expectations``, ∫ q(F) log p(Y|F) dF, needs to
be computed under a Gaussian distribution q(F) ~ N(μ, Σ).

The :func:`variational_expectations <gpflow.likelihoods.Likelihood.variational_expectations>`
method can be overriden if this can be computed in closed form, otherwise; if
the new likelihood inherits
:class:`Likelihood <gpflow.likelihoods.Likelihood>` the default will use
Gauss-Hermite numerical integration (works well when F is 1D
or 2D), if the new likelihood inherits from
:class:`MonteCarloLikelihood <gpflow.likelihoods.MonteCarloLikelihood>` the
integration is done by sampling (can be more suitable when F is higher dimensional).
"""

from .base import (
    Likelihood,
    MonteCarloLikelihood,
    QuadratureLikelihood,
    ScalarLikelihood,
    SwitchedLikelihood,
)
from .misc import GaussianMC
from .multiclass import MultiClass, RobustMax, Softmax
from .multilatent import (
    HeteroskedasticTFPConditional,
    MultiLatentLikelihood,
    MultiLatentTFPConditional,
)
from .scalar_continuous import Beta, Exponential, Gamma, Gaussian, StudentT
from .scalar_discrete import Bernoulli, Ordinal, Poisson

__all__ = [
    "Bernoulli",
    "Beta",
    "Exponential",
    "Gamma",
    "Gaussian",
    "GaussianMC",
    "HeteroskedasticTFPConditional",
    "Likelihood",
    "MonteCarloLikelihood",
    "MultiClass",
    "MultiLatentLikelihood",
    "MultiLatentTFPConditional",
    "Ordinal",
    "Poisson",
    "QuadratureLikelihood",
    "RobustMax",
    "ScalarLikelihood",
    "Softmax",
    "StudentT",
    "SwitchedLikelihood",
    "base",
    "misc",
    "multiclass",
    "multilatent",
    "scalar_continuous",
    "scalar_discrete",
    "utils",
]


# gpflow/likelihoods/base.py

# Copyright 2016-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.
import abc
from typing import Any, Callable, Iterable, Optional, Sequence, Union

import tensorflow as tf

from ..base import MeanAndVariance, Module, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..quadrature import GaussianQuadrature, NDiagGHQuadrature, ndiag_mc

DEFAULT_NUM_GAUSS_HERMITE_POINTS = 20
"""
The number of Gauss-Hermite points to use for quadrature (fallback when a
likelihood method does not have an analytic method) if quadrature object is not
explicitly passed to likelihood constructor.
"""


class Likelihood(Module, abc.ABC):
    def __init__(
        self, input_dim: Optional[int], latent_dim: Optional[int], observation_dim: Optional[int]
    ) -> None:
        """
        A base class for likelihoods, which specifies an observation model
        connecting the latent functions ('F') to the data ('Y').

        All of the members of this class are expected to obey some shape conventions, as specified
        by latent_dim and observation_dim.

        If we're operating on an array of function values 'F', then the last dimension represents
        multiple functions (preceding dimensions could represent different data points, or
        different random samples, for example). Similarly, the last dimension of Y represents a
        single data point. We check that the dimensions are as this object expects.

        The return shapes of all functions in this class is the broadcasted shape of the arguments,
        excluding the last dimension of each argument.

        :param input_dim: the dimension of the input vector X for a single data point
        :param latent_dim: the dimension of the vector F of latent functions for a single data point
        :param observation_dim: the dimension of the observation vector Y for a single data point
        """
        super().__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.observation_dim = observation_dim

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        """
        The log probability density log p(Y|X,F)

        :param X: input tensor
        :param F: function evaluation tensor
        :param Y: observation tensor
        :returns: log pdf
        """
        return self._log_prob(X, F, Y)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional mean of Y|X,F: [E[Y₁|X,F], ..., E[Yₖ|X,F]]
        where K = observation_dim

        :param X: input tensor
        :param F: function evaluation tensor
        :returns: mean
        """
        return self._conditional_mean(X, F)

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional marginal variance of Y|X,F: [var(Y₁|X,F), ..., var(Yₖ|X,F)]
        where K = observation_dim

        :param X: input tensor
        :param F: function evaluation tensor
        :returns: variance
        """
        return self._conditional_variance(X, F)

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., observation_dim]",
    )
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "return[0]: [batch..., observation_dim]",
        "return[1]: [batch..., observation_dim]",
    )
    def predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        """
        Given a Normal distribution for the latent function,
        return the mean and marginal variance of Y,

        i.e. if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes the predictive mean

           ∫∫ y p(y|f)q(f) df dy

        and the predictive variance

           ∫∫ y² p(y|f)q(f) df dy  - [ ∫∫ y p(y|f)q(f) df dy ]²

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :returns: mean and variance
        """
        return self._predict_mean_and_var(X, Fmu, Fvar)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "return[0]: [batch..., observation_dim]",
        "return[1]: [batch..., observation_dim]",
    )
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Given a Normal distribution for the latent function, and a datum Y,
        compute the log predictive density of Y,

        i.e. if
            q(F) = N(Fmu, Fvar)

        and this object represents

            p(y|F)

        then this method computes the predictive density

            log ∫ p(y=Y|F)q(F) df

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: log predictive density
        """
        return self._predict_log_density(X, Fmu, Fvar, Y)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        raise NotImplementedError

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Compute the expected log density of the data, given a Gaussian
        distribution for the function values,

        i.e. if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes

           ∫ log(p(y=Y|f)) q(f) df.

        This only works if the broadcasting dimension of the statistics of q(f) (mean and variance)
        are broadcastable with that of the data Y.

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: expected log density of the data given q(F)
        """
        return self._variational_expectations(X, Fmu, Fvar, Y)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch...]",
    )
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        raise NotImplementedError


class QuadratureLikelihood(Likelihood, abc.ABC):
    def __init__(
        self,
        input_dim: Optional[int],
        latent_dim: Optional[int],
        observation_dim: Optional[int],
        *,
        quadrature: Optional[GaussianQuadrature] = None,
    ) -> None:
        super().__init__(
            input_dim=input_dim, latent_dim=latent_dim, observation_dim=observation_dim
        )
        if quadrature is None:
            with tf.init_scope():
                quadrature = NDiagGHQuadrature(
                    self._quadrature_dim, DEFAULT_NUM_GAUSS_HERMITE_POINTS
                )
        self.quadrature = quadrature

    @property
    def _quadrature_dim(self) -> int:
        """
        This defines the number of dimensions over which to evaluate the
        quadrature. Generally, this is equal to self.latent_dim. This exists
        as a separate property to allow the ScalarLikelihood subclass to
        override it with 1 (broadcasting over observation/latent dimensions
        instead).
        """
        assert self.latent_dim is not None
        return self.latent_dim

    @check_shapes(
        "F: [broadcast batch..., latent_dim]",
        "X: [broadcast batch..., input_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch..., d]",
    )
    def _quadrature_log_prob(self, F: TensorType, X: TensorType, Y: TensorType) -> tf.Tensor:
        """
        Returns the appropriate log prob integrand for quadrature.

        Quadrature expects f(X), here logp(F), to return shape [N_quad_points]
        + batch_shape + [d']. Here d'=1, but log_prob() only returns
        [N_quad_points] + batch_shape, so we add an extra dimension.

        Also see _quadrature_reduction.
        """
        return tf.expand_dims(self.log_prob(X, F, Y), axis=-1)

    @check_shapes(
        "quadrature_result: [batch..., d]",
        "return: [batch...]",
    )
    def _quadrature_reduction(self, quadrature_result: TensorType) -> tf.Tensor:
        """
        Converts the quadrature integral appropriately.

        The return shape of quadrature is batch_shape + [d']. Here, d'=1, but
        we want predict_log_density and variational_expectations to return just
        batch_shape, so we squeeze out the extra dimension.

        Also see _quadrature_log_prob.
        """
        return tf.squeeze(quadrature_result, axis=-1)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (Gaussian, Poisson) will implement specific cases.
        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: log predictive density
        """
        return self._quadrature_reduction(
            self.quadrature.logspace(self._quadrature_log_prob, Fmu, Fvar, X=X, Y=Y)
        )

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (Gaussian, Poisson) will implement specific cases.
        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :param Y: observation tensor
        :returns: variational expectations
        """
        return self._quadrature_reduction(
            self.quadrature(self._quadrature_log_prob, Fmu, Fvar, X=X, Y=Y)
        )

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        r"""
        Here, we implement a default Gauss-Hermite quadrature routine, but some
        likelihoods (e.g. Gaussian) will implement specific cases.

        :param X: input tensor
        :param Fmu: mean function evaluation tensor
        :param Fvar: variance of function evaluation tensor
        :returns: mean and variance of Y
        """

        def conditional_mean(F: TensorType, X_: TensorType) -> tf.Tensor:
            return self.conditional_mean(X_, F)

        def conditional_y_squared(F: TensorType, X_: TensorType) -> tf.Tensor:
            return self.conditional_variance(X_, F) + tf.square(self.conditional_mean(X_, F))

        E_y, E_y2 = self.quadrature([conditional_mean, conditional_y_squared], Fmu, Fvar, X_=X)
        V_y = E_y2 - E_y ** 2
        return E_y, V_y


class ScalarLikelihood(QuadratureLikelihood, abc.ABC):
    """
    A likelihood class that helps with scalar likelihood functions: likelihoods where
    each scalar latent function is associated with a single scalar observation variable.

    If there are multiple latent functions, then there must be a corresponding number of data: we
    check for this.

    The `Likelihood` class contains methods to compute marginal statistics of functions
    of the latents and the data ϕ(y,x,f):

    * variational_expectations:  ϕ(y,x,f) = log p(y|x,f)
    * predict_log_density: ϕ(y,x,f) = p(y|x,f)

    Those statistics are computed after having first marginalized the latent processes f
    under a multivariate normal distribution q(x,f) that is fully factorized.

    Some univariate integrals can be done by quadrature: we implement quadrature routines for 1D
    integrals in this class, though they may be overwritten by inheriting classes where those
    integrals are available in closed form.
    """

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(input_dim=None, latent_dim=None, observation_dim=None, **kwargs)

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        r"""
        Compute log p(Y|X,F), where by convention we sum out the last axis as it represented
        independent latent functions and observations.
        :param F: function evaluation tensor
        :param Y: observation tensor
        """
        return tf.reduce_sum(self._scalar_log_prob(X, F, Y), axis=-1)

    @abc.abstractmethod
    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "F: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch..., latent_dim]",
    )
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @property
    def _quadrature_dim(self) -> int:
        """
        Quadrature is over the latent dimensions. Generally, this is equal to
        self.latent_dim. This separate property allows the ScalarLikelihood
        subclass to override it with 1 (broadcasting over observation/latent
        dimensions instead).
        """
        return 1

    @inherit_check_shapes
    def _quadrature_log_prob(self, F: TensorType, X: TensorType, Y: TensorType) -> tf.Tensor:
        """
        Returns the appropriate log prob integrand for quadrature.

        Quadrature expects f(X), here logp(F), to return shape [N_quad_points]
        + batch_shape + [d']. Here d' corresponds to the last dimension of both
        F and Y, and _scalar_log_prob simply broadcasts over this.

        Also see _quadrature_reduction.
        """
        return self._scalar_log_prob(X, F, Y)

    @inherit_check_shapes
    def _quadrature_reduction(self, quadrature_result: TensorType) -> tf.Tensor:
        """
        Converts the quadrature integral appropriately.

        The return shape of quadrature is batch_shape + [d']. Here, d'
        corresponds to the last dimension of both F and Y, and we want to sum
        over the observations to obtain the overall predict_log_density or
        variational_expectations.

        Also see _quadrature_log_prob.
        """
        return tf.reduce_sum(quadrature_result, axis=-1)


class SwitchedLikelihood(ScalarLikelihood):
    def __init__(self, likelihood_list: Iterable[ScalarLikelihood], **kwargs: Any) -> None:
        """
        In this likelihood, we assume at extra column of Y, which contains
        integers that specify a likelihood from the list of likelihoods.
        """
        super().__init__(**kwargs)
        self.likelihoods = list(likelihood_list)

    @check_shapes(
        "args[all]: [batch..., .]",
        "return: [batch..., ...]",
    )
    def _partition_and_stitch(self, args: Sequence[TensorType], func_name: str) -> tf.Tensor:
        """
        args is a list of tensors, to be passed to self.likelihoods.<func_name>

        args[-1] is the 'Y' argument, which contains the indexes to self.likelihoods.

        This function splits up the args using dynamic_partition, calls the
        relevant function on the likelihoods, and re-combines the result.
        """
        # get the index from Y
        args_list = list(args)
        Y = args_list[-1]
        ind = Y[..., -1]
        ind = tf.cast(ind, tf.int32)
        Y = Y[..., :-1]
        args_list[-1] = Y

        # split up the arguments into chunks corresponding to the relevant likelihoods
        args_chunks = zip(*[tf.dynamic_partition(X, ind, len(self.likelihoods)) for X in args_list])

        # apply the likelihood-function to each section of the data
        funcs = [getattr(lik, func_name) for lik in self.likelihoods]
        results = [f(*args_i) for f, args_i in zip(funcs, args_chunks)]

        # stitch the results back together
        partitions = tf.dynamic_partition(tf.range(0, tf.size(ind)), ind, len(self.likelihoods))
        results = tf.dynamic_stitch(partitions, results)

        return results

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return self._partition_and_stitch([X, F, Y], "_scalar_log_prob")

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return self._partition_and_stitch([X, Fmu, Fvar, Y], "predict_log_density")

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return self._partition_and_stitch([X, Fmu, Fvar, Y], "variational_expectations")

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        mvs = [lik.predict_mean_and_var(X, Fmu, Fvar) for lik in self.likelihoods]
        mu_list, var_list = zip(*mvs)
        mu = tf.concat(mu_list, axis=1)
        var = tf.concat(var_list, axis=1)
        return mu, var

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        raise NotImplementedError


class MonteCarloLikelihood(Likelihood):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.num_monte_carlo_points = 100

    @check_shapes(
        "Fmu: [batch..., latent_dim]",
        "Fvar: [batch..., latent_dim]",
        "Ys.values(): [batch..., .]",
        "return: [broadcast n_funcs, batch..., .]",
    )
    def _mc_quadrature(
        self,
        funcs: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
        Fmu: TensorType,
        Fvar: TensorType,
        logspace: bool = False,
        epsilon: Optional[TensorType] = None,
        **Ys: TensorType,
    ) -> tf.Tensor:
        return ndiag_mc(funcs, self.num_monte_carlo_points, Fmu, Fvar, logspace, epsilon, **Ys)

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, epsilon: Optional[TensorType] = None
    ) -> MeanAndVariance:
        r"""
        Given a Normal distribution for the latent function,
        return the mean of Y

        if
            q(f) = N(Fmu, Fvar)

        and this object represents

            p(y|f)

        then this method computes the predictive mean

           ∫∫ y p(y|f)q(f) df dy

        and the predictive variance

           ∫∫ y² p(y|f)q(f) df dy  - [ ∫∫ y p(y|f)q(f) df dy ]²

        Here, we implement a default Monte Carlo routine.
        """

        def conditional_mean(F: TensorType, X_: TensorType) -> TensorType:
            return self.conditional_mean(X_, F)

        def conditional_y_squared(F: TensorType, X_: TensorType) -> TensorType:
            return self.conditional_variance(X_, F) + tf.square(self.conditional_mean(X_, F))

        E_y, E_y2 = self._mc_quadrature(
            [conditional_mean, conditional_y_squared],
            Fmu,
            Fvar,
            epsilon=epsilon,
            X_=X,
        )
        V_y = E_y2 - tf.square(E_y)
        return E_y, V_y  # [N, D]

    @inherit_check_shapes
    def _predict_log_density(
        self,
        X: TensorType,
        Fmu: TensorType,
        Fvar: TensorType,
        Y: TensorType,
        epsilon: Optional[TensorType] = None,
    ) -> tf.Tensor:
        r"""
        Given a Normal distribution for the latent function, and a datum Y,
        compute the log predictive density of Y.

        i.e. if
            q(x, f) = N(Fmu, Fvar)

        and this object represents

            p(y|x,f)

        then this method computes the predictive density

            log ∫ p(y=Y|x,f)q(x,f) df

        Here, we implement a default Monte Carlo routine.
        """

        def log_prob(F: TensorType, X_: TensorType, Y_: TensorType) -> tf.Tensor:
            return self.log_prob(X_, F, Y_)

        return tf.reduce_sum(
            self._mc_quadrature(log_prob, Fmu, Fvar, logspace=True, epsilon=epsilon, X_=X, Y_=Y),
            axis=-1,
        )

    @inherit_check_shapes
    def _variational_expectations(
        self,
        X: TensorType,
        Fmu: TensorType,
        Fvar: TensorType,
        Y: TensorType,
        epsilon: Optional[TensorType] = None,
    ) -> tf.Tensor:
        r"""
        Compute the expected log density of the data, given a Gaussian
        distribution for the function values.

        if
            q(x,f) = N(Fmu, Fvar)  - Fmu: [N, D]  Fvar: [N, D]

        and this object represents

            p(y|x,f)  - Y: [N, 1]

        then this method computes

           ∫ (log p(y|x,f)) q(x,f) df.


        Here, we implement a default Monte Carlo quadrature routine.
        """

        def log_prob(F: TensorType, X_: TensorType, Y_: TensorType) -> tf.Tensor:
            return self.log_prob(X_, F, Y_)

        return tf.reduce_sum(
            self._mc_quadrature(log_prob, Fmu, Fvar, epsilon=epsilon, X_=X, Y_=Y), axis=-1
        )


# gpflow/likelihoods/misc.py

# 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 .base import MonteCarloLikelihood
from .scalar_continuous import Gaussian


class GaussianMC(MonteCarloLikelihood, Gaussian):  # type: ignore[misc]
    """
    Stochastic version of Gaussian likelihood for demonstration purposes only.
    """

    pass


# gpflow/likelihoods/multiclass.py

# 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 typing import Any, Optional

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from ..base import MeanAndVariance, Module, Parameter, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..quadrature import hermgauss
from ..utilities import to_default_float, to_default_int
from .base import Likelihood, MonteCarloLikelihood


class Softmax(MonteCarloLikelihood):
    """
    The soft-max multi-class likelihood.  It can only provide a stochastic
    Monte-Carlo estimate of the variational expectations term, but this
    added variance tends to be small compared to that due to mini-batching
    (when using the SVGP model).
    """

    def __init__(self, num_classes: int, **kwargs: Any) -> None:
        super().__init__(input_dim=None, latent_dim=num_classes, observation_dim=None, **kwargs)
        self.num_classes = self.latent_dim

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return -tf.nn.sparse_softmax_cross_entropy_with_logits(logits=F, labels=Y[:, 0])

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return tf.nn.softmax(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        p = self.conditional_mean(X, F)
        return p - p ** 2


class RobustMax(Module):
    r"""
    This class represent a multi-class inverse-link function. Given a vector
    :math:`f=[f_1, f_2, ... f_k]`, the result of the mapping is

    .. math::

       y = [y_1 ... y_k]

    with

    .. math::

       y_i = \left\{
       \begin{array}{ll}
           (1-\varepsilon)   & \textrm{if} \ i = \textrm{argmax}(f) \\
           \varepsilon/(k-1) & \textrm{otherwise}
       \end{array}
       \right.

    where :math:`k` is the number of classes.
    """

    @check_shapes(
        "epsilon: []",
    )
    def __init__(self, num_classes: int, epsilon: float = 1e-3, **kwargs: Any) -> None:
        """
        `epsilon` represents the fraction of 'errors' in the labels of the
        dataset. This may be a hard parameter to optimize, so by default
        it is set un-trainable, at a small value.
        """
        super().__init__(**kwargs)
        transform = tfp.bijectors.Sigmoid()
        prior = tfp.distributions.Beta(to_default_float(0.2), to_default_float(5.0))
        self.epsilon = Parameter(epsilon, transform=transform, prior=prior, trainable=False)
        self.num_classes = num_classes
        self._squash = 1e-6

    @check_shapes(
        "F: [broadcast batch..., latent_dim]",
        "return: [batch..., latent_dim]",
    )
    def __call__(self, F: TensorType) -> tf.Tensor:
        i = tf.argmax(F, 1)
        return tf.one_hot(i, self.num_classes, 1.0 - self.epsilon, self.eps_k1)

    @property  # type: ignore[misc]  # Mypy doesn't like decorated properties.
    @check_shapes(
        "return: []",
    )
    def eps_k1(self) -> tf.Tensor:
        return self.epsilon / (self.num_classes - 1.0)

    @check_shapes(
        "val: [batch...]",
        "return: [batch...]",
    )
    def safe_sqrt(self, val: TensorType) -> tf.Tensor:
        return tf.sqrt(tf.maximum(val, 1e-10))

    @check_shapes(
        "Y: [broadcast batch..., observation_dim]",
        "mu: [broadcast batch..., latent_dim]",
        "var: [broadcast batch..., latent_dim]",
        "gh_x: [n_quad_points]",
        "gh_w: [n_quad_points]",
        "return: [batch..., observation_dim]",
    )
    def prob_is_largest(
        self, Y: TensorType, mu: TensorType, var: TensorType, gh_x: TensorType, gh_w: TensorType
    ) -> tf.Tensor:
        Y = to_default_int(Y)
        # work out what the mean and variance is of the indicated latent function.
        oh_on = tf.cast(
            tf.one_hot(tf.reshape(Y, (-1,)), self.num_classes, 1.0, 0.0), dtype=mu.dtype
        )
        mu_selected = tf.reduce_sum(oh_on * mu, 1)
        var_selected = tf.reduce_sum(oh_on * var, 1)

        # generate Gauss Hermite grid
        X = tf.reshape(mu_selected, (-1, 1)) + gh_x * tf.reshape(
            self.safe_sqrt(2.0 * var_selected), (-1, 1)
        )

        # compute the CDF of the Gaussian between the latent functions and the grid (including the selected function)
        dist = (tf.expand_dims(X, 1) - tf.expand_dims(mu, 2)) / tf.expand_dims(
            self.safe_sqrt(var), 2
        )
        cdfs = 0.5 * (1.0 + tf.math.erf(dist / np.sqrt(2.0)))

        cdfs = cdfs * (1 - 2 * self._squash) + self._squash

        # blank out all the distances on the selected latent function
        oh_off = tf.cast(
            tf.one_hot(tf.reshape(Y, (-1,)), self.num_classes, 0.0, 1.0), dtype=mu.dtype
        )
        cdfs = cdfs * tf.expand_dims(oh_off, 2) + tf.expand_dims(oh_on, 2)

        # take the product over the latent functions, and the sum over the GH grid.
        return tf.reduce_prod(cdfs, axis=[1]) @ tf.reshape(gh_w / np.sqrt(np.pi), (-1, 1))


class MultiClass(Likelihood):
    def __init__(
        self, num_classes: int, invlink: Optional[RobustMax] = None, **kwargs: Any
    ) -> None:
        """
        A likelihood for multi-way classification.  Currently the only valid
        choice of inverse-link function (invlink) is an instance of RobustMax.

        For most problems, the stochastic `Softmax` likelihood may be more
        appropriate (note that you then cannot use Scipy optimizer).
        """
        super().__init__(input_dim=None, latent_dim=num_classes, observation_dim=None, **kwargs)
        self.num_classes = num_classes
        self.num_gauss_hermite_points = 20

        if invlink is None:
            invlink = RobustMax(self.num_classes)

        if not isinstance(invlink, RobustMax):
            raise NotImplementedError

        self.invlink = invlink

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        hits = tf.equal(tf.expand_dims(tf.argmax(F, 1), 1), tf.cast(Y, tf.int64))
        yes = tf.ones(tf.shape(Y), dtype=default_float()) - self.invlink.epsilon
        no = tf.zeros(tf.shape(Y), dtype=default_float()) + self.invlink.eps_k1
        p = tf.where(hits, yes, no)
        return tf.reduce_sum(tf.math.log(p), axis=-1)

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        gh_x, gh_w = hermgauss(self.num_gauss_hermite_points)
        p = self.invlink.prob_is_largest(Y, Fmu, Fvar, gh_x, gh_w)
        ve = p * tf.math.log(1.0 - self.invlink.epsilon) + (1.0 - p) * tf.math.log(
            self.invlink.eps_k1
        )
        return tf.reduce_sum(ve, axis=-1)

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        possible_outputs = [
            tf.fill(tf.stack([tf.shape(Fmu)[0], 1]), np.array(i, dtype=np.int64))
            for i in range(self.num_classes)
        ]
        ps = [self._predict_non_logged_density(X, Fmu, Fvar, po) for po in possible_outputs]
        ps = tf.transpose(tf.stack([tf.reshape(p, (-1,)) for p in ps]))
        return ps, ps - tf.square(ps)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return tf.reduce_sum(
            tf.math.log(self._predict_non_logged_density(X, Fmu, Fvar, Y)), axis=-1
        )

    @check_shapes(
        "X: [broadcast batch..., input_dim]",
        "Fmu: [broadcast batch..., latent_dim]",
        "Fvar: [broadcast batch..., latent_dim]",
        "Y: [broadcast batch..., observation_dim]",
        "return: [batch..., observation_dim]",
    )
    def _predict_non_logged_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        gh_x, gh_w = hermgauss(self.num_gauss_hermite_points)
        p = self.invlink.prob_is_largest(Y, Fmu, Fvar, gh_x, gh_w)
        den = p * (1.0 - self.invlink.epsilon) + (1.0 - p) * (self.invlink.eps_k1)
        return den

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        p = self.conditional_mean(X, F)
        return p - tf.square(p)


# gpflow/likelihoods/multilatent.py

# Copyright 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 typing import Any, Callable, Optional, Type

import tensorflow as tf
import tensorflow_probability as tfp

from ..base import TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive
from .base import QuadratureLikelihood


class MultiLatentLikelihood(QuadratureLikelihood):
    r"""
    A Likelihood which assumes that a single dimensional observation is driven
    by multiple latent GPs.

    Note that this implementation does not allow for taking into account
    covariance between outputs.
    """

    def __init__(self, latent_dim: int, **kwargs: Any) -> None:
        super().__init__(
            input_dim=None,
            latent_dim=latent_dim,
            observation_dim=1,
            **kwargs,
        )


class MultiLatentTFPConditional(MultiLatentLikelihood):
    """
    MultiLatent likelihood where the conditional distribution
    is given by a TensorFlow Probability Distribution.
    """

    def __init__(
        self,
        latent_dim: int,
        conditional_distribution: Callable[..., tfp.distributions.Distribution],
        **kwargs: Any,
    ):
        """
        :param latent_dim: number of arguments to the `conditional_distribution` callable
        :param conditional_distribution: function from F to a tfp Distribution,
            where F has shape [..., latent_dim]
        """
        super().__init__(latent_dim, **kwargs)
        self.conditional_distribution = conditional_distribution

    @inherit_check_shapes
    def _log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        """
        The log probability density log p(Y|F)

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :param Y: observation Tensor, with shape [..., 1]:
        :returns: log pdf, with shape [...]
        """
        return tf.squeeze(self.conditional_distribution(F).log_prob(Y), -1)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional marginal mean of Y|F: [E(Y₁|F)]

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :returns: mean [..., 1]
        """
        return self.conditional_distribution(F).mean()

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        """
        The conditional marginal variance of Y|F: [Var(Y₁|F)]

        :param F: function evaluation Tensor, with shape [..., latent_dim]
        :returns: variance [..., 1]
        """
        return self.conditional_distribution(F).variance()


class HeteroskedasticTFPConditional(MultiLatentTFPConditional):
    """
    Heteroskedastic Likelihood where the conditional distribution
    is given by a TensorFlow Probability Distribution.
    The `loc` and `scale` of the distribution are given by a
    two-dimensional multi-output GP.
    """

    def __init__(
        self,
        distribution_class: Type[tfp.distributions.Distribution] = tfp.distributions.Normal,
        scale_transform: Optional[tfp.bijectors.Bijector] = None,
        **kwargs: Any,
    ) -> None:
        """
        :param distribution_class: distribution class parameterized by `loc` and `scale`
            as first and second argument, respectively.
        :param scale_transform: callable/bijector applied to the latent
            function modelling the scale to ensure its positivity.
            Typically, `tf.exp` or `tf.softplus`, but can be any function f: R -> R^+. Defaults to exp if not explicitly specified.
        """
        if scale_transform is None:
            scale_transform = positive(base="exp")
        self.scale_transform = scale_transform

        @check_shapes(
            "F: [batch..., 2]",
        )
        def conditional_distribution(F: TensorType) -> tfp.distributions.Distribution:
            loc = F[..., :1]
            scale = self.scale_transform(F[..., 1:])
            return distribution_class(loc, scale)

        super().__init__(
            latent_dim=2,
            conditional_distribution=conditional_distribution,
            **kwargs,
        )


# gpflow/likelihoods/scalar_continuous.py

# 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 math import sqrt
from typing import Any, Callable, Optional

import numpy as np
import tensorflow as tf

from .. import logdensities
from ..base import MeanAndVariance, TensorType
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities.parameter_or_function import (
    ConstantOrFunction,
    ParameterOrFunction,
    evaluate_parameter_or_function,
    prepare_parameter_or_function,
)
from .base import ScalarLikelihood
from .utils import inv_probit

DEFAULT_LOWER_BOUND = 1e-6


class Gaussian(ScalarLikelihood):
    r"""
    The Gaussian likelihood is appropriate where uncertainties associated with
    the data are believed to follow a normal distribution, with constant
    variance.

    Very small uncertainties can lead to numerical instability during the
    optimization process. A lower bound of 1e-6 is therefore imposed on the
    likelihood variance by default.
    """

    def __init__(
        self,
        variance: Optional[ConstantOrFunction] = None,
        *,
        scale: Optional[ConstantOrFunction] = None,
        variance_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        """
        :param variance: The noise variance; must be greater than
            ``variance_lower_bound``. This is mutually exclusive with `scale`.
        :param scale: The noise scale; must be greater than
            ``sqrt(variance_lower_bound)``. This is mutually exclusive with `variance`.
        :param variance_lower_bound: The lower (exclusive) bound of ``variance``.
        :param kwargs: Keyword arguments forwarded to :class:`ScalarLikelihood`.
        """
        super().__init__(**kwargs)

        self.variance_lower_bound = variance_lower_bound
        self.scale_lower_bound = sqrt(variance_lower_bound)
        if scale is None:
            if variance is None:
                variance = 1.0
            self.variance: Optional[ParameterOrFunction] = prepare_parameter_or_function(
                variance, lower_bound=self.variance_lower_bound
            )
            self.scale: Optional[ParameterOrFunction] = None
        else:
            if variance is None:
                self.variance = None
                self.scale = prepare_parameter_or_function(
                    scale, lower_bound=self.scale_lower_bound
                )
            else:
                assert False, "Cannot set both `variance` and `scale`."

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _variance(self, X: TensorType) -> tf.Tensor:
        if self.variance is not None:
            return evaluate_parameter_or_function(
                self.variance, X, lower_bound=self.variance_lower_bound
            )
        else:
            assert self.scale is not None  # For mypy.
            return (
                evaluate_parameter_or_function(self.scale, X, lower_bound=self.scale_lower_bound)
                ** 2
            )

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [batch..., N, 1]",
    )
    def variance_at(self, X: TensorType) -> tf.Tensor:
        variance = self._variance(X)
        shape = tf.concat([tf.shape(X)[:-1], [1]], 0)
        return tf.broadcast_to(variance, shape)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.gaussian(Y, F, self._variance(X))

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:  # pylint: disable=R0201
        return tf.identity(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        shape = tf.shape(F)
        return tf.broadcast_to(self._variance(X), shape)

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        return tf.identity(Fmu), Fvar + self._variance(X)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        return tf.reduce_sum(logdensities.gaussian(Y, Fmu, Fvar + self._variance(X)), axis=-1)

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        variance = self._variance(X)
        return tf.reduce_sum(
            -0.5 * np.log(2 * np.pi)
            - 0.5 * tf.math.log(variance)
            - 0.5 * ((Y - Fmu) ** 2 + Fvar) / variance,
            axis=-1,
        )


class Exponential(ScalarLikelihood):
    def __init__(self, invlink: Callable[[tf.Tensor], tf.Tensor] = tf.exp, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.exponential(Y, self.invlink(F))

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return tf.square(self.invlink(F))

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        if self.invlink is tf.exp:
            return tf.reduce_sum(-tf.exp(-Fmu + Fvar / 2) * Y - Fmu, axis=-1)
        return super()._variational_expectations(X, Fmu, Fvar, Y)


class StudentT(ScalarLikelihood):
    def __init__(
        self,
        scale: ConstantOrFunction = 1.0,
        df: float = 3.0,
        scale_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        """
        :param scale float: scale parameter
        :param df float: degrees of freedom
        """
        super().__init__(**kwargs)
        self.df = df
        self.scale_lower_bound = scale_lower_bound
        self.scale = prepare_parameter_or_function(scale, lower_bound=self.scale_lower_bound)

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _scale(self, X: TensorType) -> tf.Tensor:
        return evaluate_parameter_or_function(self.scale, X, lower_bound=self.scale_lower_bound)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.student_t(Y, F, self._scale(X), self.df)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return F

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        shape = tf.shape(F)
        var = (self._scale(X) ** 2) * (self.df / (self.df - 2.0))
        return tf.broadcast_to(var, shape)


class Gamma(ScalarLikelihood):
    """
    Use the transformed GP to give the *scale* (inverse rate) of the Gamma
    """

    def __init__(
        self,
        invlink: Callable[[tf.Tensor], tf.Tensor] = tf.exp,
        shape: ConstantOrFunction = 1.0,
        shape_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink
        self.shape_lower_bound = shape_lower_bound
        self.shape = prepare_parameter_or_function(shape, lower_bound=self.shape_lower_bound)

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _shape(self, X: TensorType) -> tf.Tensor:
        return evaluate_parameter_or_function(self.shape, X, lower_bound=self.shape_lower_bound)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.gamma(Y, self._shape(X), self.invlink(F))

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self._shape(X) * self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        scale = self.invlink(F)
        return self._shape(X) * (scale ** 2)

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        if self.invlink is tf.exp:
            shape = self._shape(X)
            return tf.reduce_sum(
                -shape * Fmu
                - tf.math.lgamma(shape)
                + (shape - 1.0) * tf.math.log(Y)
                - Y * tf.exp(-Fmu + Fvar / 2.0),
                axis=-1,
            )
        else:
            return super()._variational_expectations(X, Fmu, Fvar, Y)


class Beta(ScalarLikelihood):
    """
    This uses a reparameterisation of the Beta density. We have the mean of the
    Beta distribution given by the transformed process:

        m = invlink(f)

    and a scale parameter. The familiar α, β parameters are given by

        m     = α / (α + β)
        scale = α + β

    so:
        α = scale * m
        β  = scale * (1-m)
    """

    def __init__(
        self,
        invlink: Callable[[tf.Tensor], tf.Tensor] = inv_probit,
        scale: ConstantOrFunction = 1.0,
        scale_lower_bound: float = DEFAULT_LOWER_BOUND,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.scale_lower_bound = DEFAULT_LOWER_BOUND
        self.scale = prepare_parameter_or_function(scale, lower_bound=self.scale_lower_bound)
        self.invlink = invlink

    @check_shapes(
        "X: [batch..., N, D]",
        "return: [broadcast batch..., broadcast N, broadcast P]",
    )
    def _scale(self, X: TensorType) -> tf.Tensor:
        return evaluate_parameter_or_function(self.scale, X, lower_bound=self.scale_lower_bound)

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        mean = self.invlink(F)
        scale = self._scale(X)
        alpha = mean * scale
        beta = scale - alpha
        return logdensities.beta(Y, alpha, beta)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        mean = self.invlink(F)
        var = (mean - tf.square(mean)) / (self._scale(X) + 1.0)
        shape = tf.shape(F)
        return tf.broadcast_to(var, shape)


# gpflow/likelihoods/scalar_discrete.py

# 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 typing import Any, Callable

import numpy as np
import tensorflow as tf

from .. import logdensities
from ..base import AnyNDArray, MeanAndVariance, Parameter, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import positive, to_default_int
from .base import ScalarLikelihood
from .utils import inv_probit


class Poisson(ScalarLikelihood):
    r"""
    Poisson likelihood for use with count data, where the rate is given by the (transformed) GP.

    let g(.) be the inverse-link function, then this likelihood represents

    p(yᵢ | fᵢ) = Poisson(yᵢ | g(fᵢ) * binsize)

    Note:binsize
    For use in a Log Gaussian Cox process (doubly stochastic model) where the
    rate function of an inhomogeneous Poisson process is given by a GP.  The
    intractable likelihood can be approximated via a Riemann sum (with bins
    of size 'binsize') and using this Poisson likelihood.
    """

    def __init__(
        self,
        invlink: Callable[[tf.Tensor], tf.Tensor] = tf.exp,
        binsize: float = 1.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink
        self.binsize: AnyNDArray = np.array(binsize, dtype=default_float())

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.poisson(Y, self.invlink(F) * self.binsize)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F) * self.binsize

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F) * self.binsize

    @inherit_check_shapes
    def _variational_expectations(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        if self.invlink is tf.exp:
            return tf.reduce_sum(
                Y * Fmu
                - tf.exp(Fmu + Fvar / 2) * self.binsize
                - tf.math.lgamma(Y + 1)
                + Y * tf.math.log(self.binsize),
                axis=-1,
            )
        return super()._variational_expectations(X, Fmu, Fvar, Y)


class Bernoulli(ScalarLikelihood):
    def __init__(
        self, invlink: Callable[[tf.Tensor], tf.Tensor] = inv_probit, **kwargs: Any
    ) -> None:
        super().__init__(**kwargs)
        self.invlink = invlink

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        return logdensities.bernoulli(Y, self.invlink(F))

    @inherit_check_shapes
    def _predict_mean_and_var(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType
    ) -> MeanAndVariance:
        if self.invlink is inv_probit:
            p = inv_probit(Fmu / tf.sqrt(1 + Fvar))
            return p, p - tf.square(p)
        else:
            # for other invlink, use quadrature
            return super()._predict_mean_and_var(X, Fmu, Fvar)

    @inherit_check_shapes
    def _predict_log_density(
        self, X: TensorType, Fmu: TensorType, Fvar: TensorType, Y: TensorType
    ) -> tf.Tensor:
        p = self.predict_mean_and_var(X, Fmu, Fvar)[0]
        return tf.reduce_sum(logdensities.bernoulli(Y, p), axis=-1)

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        return self.invlink(F)

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        p = self.conditional_mean(X, F)
        return p - (p ** 2)


class Ordinal(ScalarLikelihood):
    """
    A likelihood for doing ordinal regression.

    The data are integer values from 0 to k, and the user must specify (k-1)
    'bin edges' which define the points at which the labels switch. Let the bin
    edges be [a₀, a₁, ... aₖ₋₁], then the likelihood is

    p(Y=0|F) = ɸ((a₀ - F) / σ)
    p(Y=1|F) = ɸ((a₁ - F) / σ) - ɸ((a₀ - F) / σ)
    p(Y=2|F) = ɸ((a₂ - F) / σ) - ɸ((a₁ - F) / σ)
    ...
    p(Y=K|F) = 1 - ɸ((aₖ₋₁ - F) / σ)

    where ɸ is the cumulative density function of a Gaussian (the inverse probit
    function) and σ is a parameter to be learned.

    A reference is :cite:t:`chu2005gaussian`.
    """

    @check_shapes(
        "bin_edges: [num_bins_minus_1]",
    )
    def __init__(self, bin_edges: AnyNDArray, **kwargs: Any) -> None:
        """
        bin_edges is a numpy array specifying at which function value the
        output label should switch. If the possible Y values are 0...K, then
        the size of bin_edges should be (K-1).
        """
        super().__init__(**kwargs)
        self.bin_edges = bin_edges
        self.num_bins = bin_edges.size + 1
        self.sigma = Parameter(1.0, transform=positive())

    @inherit_check_shapes
    def _scalar_log_prob(self, X: TensorType, F: TensorType, Y: TensorType) -> tf.Tensor:
        Y = to_default_int(Y)
        scaled_bins_left = tf.concat([self.bin_edges / self.sigma, np.array([np.inf])], 0)
        scaled_bins_right = tf.concat([np.array([-np.inf]), self.bin_edges / self.sigma], 0)
        selected_bins_left = tf.gather(scaled_bins_left, Y)
        selected_bins_right = tf.gather(scaled_bins_right, Y)

        return tf.math.log(
            inv_probit(selected_bins_left - F / self.sigma)
            - inv_probit(selected_bins_right - F / self.sigma)
            + 1e-6
        )

    @check_shapes(
        "F: [batch..., latent_dim]",
        "return: [batch_and_latent_dim, num_bins]",
    )
    def _make_phi(self, F: TensorType) -> tf.Tensor:
        """
        A helper function for making predictions. Constructs a probability
        matrix where each row output the probability of the corresponding
        label, and the rows match the entries of F.

        Note that a matrix of F values is flattened.
        """
        scaled_bins_left = tf.concat([self.bin_edges / self.sigma, np.array([np.inf])], 0)
        scaled_bins_right = tf.concat([np.array([-np.inf]), self.bin_edges / self.sigma], 0)
        return inv_probit(scaled_bins_left - tf.reshape(F, (-1, 1)) / self.sigma) - inv_probit(
            scaled_bins_right - tf.reshape(F, (-1, 1)) / self.sigma
        )

    @inherit_check_shapes
    def _conditional_mean(self, X: TensorType, F: TensorType) -> tf.Tensor:
        phi = self._make_phi(F)
        Ys = tf.reshape(np.arange(self.num_bins, dtype=default_float()), (-1, 1))
        return tf.reshape(tf.linalg.matmul(phi, Ys), tf.shape(F))

    @inherit_check_shapes
    def _conditional_variance(self, X: TensorType, F: TensorType) -> tf.Tensor:
        phi = self._make_phi(F)
        Ys = tf.reshape(np.arange(self.num_bins, dtype=default_float()), (-1, 1))
        E_y = phi @ Ys
        E_y2 = phi @ (Ys ** 2)
        return tf.reshape(E_y2 - E_y ** 2, tf.shape(F))


# gpflow/likelihoods/utils.py

# 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.

import numpy as np
import tensorflow as tf

from ..experimental.check_shapes import check_shapes


@check_shapes(
    "x: [batch...]",
    "return: [batch...]",
)
def inv_probit(x: tf.Tensor) -> tf.Tensor:
    jitter = 1e-3  # ensures output is strictly between 0 and 1
    return 0.5 * (1.0 + tf.math.erf(x / np.sqrt(2.0))) * (1 - 2 * jitter) + jitter


# gpflow/models/__init__.py

from .cglb import CGLB
from .gplvm import GPLVM, BayesianGPLVM
from .gpmc import GPMC
from .gpr import GPR
from .model import BayesianModel, GPModel
from .sgpmc import SGPMC
from .sgpr import GPRFITC, SGPR
from .svgp import SVGP
from .training_mixins import ExternalDataTrainingLossMixin, InternalDataTrainingLossMixin
from .util import maximum_log_likelihood_objective, training_loss, training_loss_closure
from .vgp import VGP, VGPOpperArchambeau

__all__ = [
    "BayesianGPLVM",
    "BayesianModel",
    "CGLB",
    "ExternalDataTrainingLossMixin",
    "GPLVM",
    "GPMC",
    "GPModel",
    "GPR",
    "GPRFITC",
    "InternalDataTrainingLossMixin",
    "SGPMC",
    "SGPR",
    "SVGP",
    "VGP",
    "VGPOpperArchambeau",
    "cglb",
    "gplvm",
    "gpmc",
    "gpr",
    "maximum_log_likelihood_objective",
    "model",
    "sgpmc",
    "sgpr",
    "svgp",
    "training_loss",
    "training_loss_closure",
    "training_mixins",
    "util",
    "vgp",
]


# gpflow/models/cglb.py

# Copyright 2021 the GPflow authors.
#
# 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 typing import Any, List, NamedTuple, Optional, Tuple

import tensorflow as tf

from ..base import InputData, MeanAndVariance, Parameter, RegressionData, TensorType
from ..config import default_float, default_int
from ..covariances import Kuf
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..utilities import add_noise_cov, assert_params_false, to_default_float
from .sgpr import SGPR


class CGLB(SGPR):
    """
    Conjugate Gradient Lower Bound.

    The key reference is :cite:t:`pmlr-v139-artemev21a`.

    :param cg_tolerance: Determines accuracy to which conjugate
        gradient is run when evaluating the elbo. Running more
        iterations of CG would increase the ELBO by at most
        `cg_tolerance`.

    :param max_cg_iters: Maximum number of iterations of CG to run
        per evaluation of the ELBO (or mean prediction).

    :param restart_cg_iters: How frequently to restart the CG iteration.
        Can be useful to avoid build up of numerical errors when
        many steps of CG are run.

    :param v_grad_optimization: If False, in every evaluation of the
        ELBO, CG is run to select a new auxilary vector `v`. If
        False, no CG is run when evaluating the ELBO but
        gradients with respect to `v` are tracked so that it can
        be optimized jointly with other parameters.

    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        *args: Any,
        cg_tolerance: float = 1.0,
        max_cg_iters: int = 100,
        restart_cg_iters: int = 40,
        v_grad_optimization: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(data, *args, **kwargs)
        n, b = self.data[1].shape
        self._v = Parameter(tf.zeros((b, n), dtype=default_float()), trainable=v_grad_optimization)
        self._cg_tolerance = cg_tolerance
        self._max_cg_iters = max_cg_iters
        self._restart_cg_iters = restart_cg_iters

    @property  # type: ignore[misc]  # Mypy doesn't like decorated properties.
    @check_shapes(
        "return: [N, P]",
    )
    def aux_vec(self) -> Parameter:
        return self._v

    @inherit_check_shapes
    def logdet_term(self, common: SGPR.CommonTensors) -> tf.Tensor:
        r"""
        Compute a lower bound on :math:`-0.5 * \log |K + σ²I|` based on a
        low-rank approximation to K.

        .. math::

           \log |K + σ²I| <= \log |Q + σ²I| + n * \log(1 + \textrm{tr}(K - Q)/(σ²n)).

        This bound is at least as tight as

        .. math::

           \log |K + σ²I| <=  \log |Q + σ²I| + \textrm{tr}(K - Q)/σ²,

        which appears in SGPR.
        """
        LB = common.LB
        AAT = common.AAT
        x, y = self.data
        num_data = to_default_float(tf.shape(y)[0])
        output_dim = to_default_float(tf.shape(y)[1])
        sigma_sq = self.likelihood.variance

        kdiag = self.kernel(x, full_cov=False)
        # t / σ²
        trace = tf.reduce_sum(kdiag) / sigma_sq - tf.reduce_sum(tf.linalg.diag_part(AAT))
        logdet_b = tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))
        logsigma_sq = num_data * tf.math.log(sigma_sq)

        # Correction term from Jensen's inequality
        logtrace = num_data * tf.math.log(1 + trace / num_data)
        return -output_dim * (logdet_b + 0.5 * logsigma_sq + 0.5 * logtrace)

    @inherit_check_shapes
    def quad_term(self, common: SGPR.CommonTensors) -> tf.Tensor:
        """
        Computes a lower bound on the quadratic term in the log
        marginal likelihood of conjugate GPR. The bound is based on
        an auxiliary vector, v. For :math:`Q ≺ K` and :math:`r=y - Kv`

        .. math::
            -0.5 * (rᵀQ⁻¹r + 2yᵀv - vᵀ K v ) <= -0.5 * yᵀK⁻¹y <= -0.5 * (2yᵀv - vᵀKv).

        Equality holds if :math:`r=0`, i.e. :math:`v = K⁻¹y`.

        If `self.aux_vec` is trainable, gradients are computed with
        respect to :math:`v` as well and :math:`v` can be optimized
        using gradient based methods.

        Otherwise, :math:`v` is updated with the method of conjugate
        gradients (CG). CG is run until :math:`0.5 * rᵀQ⁻¹r <= ϵ`,
        which ensures that the maximum bias due to this term is not
        more than :math:`ϵ`. The :math:`ϵ` is the CG tolerance.
        """
        x, y = self.data
        err = y - self.mean_function(x)
        sigma_sq = self.likelihood.variance
        K = add_noise_cov(self.kernel.K(x), sigma_sq)

        A = common.A
        LB = common.LB
        preconditioner = NystromPreconditioner(A, LB, sigma_sq)
        err_t = tf.transpose(err)

        v_init = self.aux_vec
        if not v_init.trainable:
            v = cglb_conjugate_gradient(
                K,
                err_t,
                v_init,
                preconditioner,
                self._cg_tolerance,
                self._max_cg_iters,
                self._restart_cg_iters,
            )
        else:
            v = v_init

        Kv = v @ K
        r = err_t - Kv
        _, error_bound = preconditioner(r)
        lb = tf.reduce_sum(v * (r + 0.5 * Kv))
        ub = lb + 0.5 * error_bound

        if not v_init.trainable:
            v_init.assign(v)

        return -ub

    @inherit_check_shapes
    def predict_f(
        self,
        Xnew: InputData,
        full_cov: bool = False,
        full_output_cov: bool = False,
        cg_tolerance: Optional[float] = 1e-3,
    ) -> MeanAndVariance:
        """
        The posterior mean for CGLB model is given by

        .. :math::
            m(xs) = K_{sf}v + Q_{ff}Q⁻¹r

        where :math:`r = y - K v`  is the residual from CG.

        Note that when :math:`v=0`, this agree with the SGPR mean,
        while if :math:`v = K⁻¹ y`, then :math:`r=0`, and the exact
        GP mean is recovered.

        :param cg_tolerance: float or None: If None, the cached value of
            :math:`v` is used. If float, conjugate gradient is run until :math:`rᵀQ⁻¹r < ϵ`.
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        x, y = self.data
        err = y - self.mean_function(x)
        kxx = self.kernel(x, x)
        ksf = self.kernel(Xnew, x)
        sigma_sq = self.likelihood.variance
        sigma = tf.sqrt(sigma_sq)
        iv = self.inducing_variable
        kernel = self.kernel
        matmul = tf.linalg.matmul
        trisolve = tf.linalg.triangular_solve

        kmat = add_noise_cov(kxx, sigma_sq)

        common = self._common_calculation()
        A, LB, L = common.A, common.LB, common.L

        v = self.aux_vec
        if cg_tolerance is not None:
            preconditioner = NystromPreconditioner(A, LB, sigma_sq)
            err_t = tf.transpose(err)
            v = cglb_conjugate_gradient(
                kmat,
                err_t,
                v,
                preconditioner,
                cg_tolerance,
                self._max_cg_iters,
                self._restart_cg_iters,
            )

        cg_mean = matmul(ksf, v, transpose_b=True)
        res = err - matmul(kmat, v, transpose_b=True)

        Kus = Kuf(iv, kernel, Xnew)
        Ares = matmul(A, res)  # The god of war!
        c = trisolve(LB, Ares, lower=True) / sigma
        tmp1 = trisolve(L, Kus, lower=True)
        tmp2 = trisolve(LB, tmp1, lower=True)
        sgpr_mean = matmul(tmp2, c, transpose_a=True)

        if full_cov:
            var = (
                kernel(Xnew)
                + matmul(tmp2, tmp2, transpose_a=True)
                - matmul(tmp1, tmp1, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), 0)
                - tf.reduce_sum(tf.square(tmp1), 0)
            )
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        mean = sgpr_mean + cg_mean + self.mean_function(Xnew)
        return mean, var

    @inherit_check_shapes
    def predict_y(
        self,
        Xnew: InputData,
        full_cov: bool = False,
        full_output_cov: bool = False,
        cg_tolerance: Optional[float] = 1e-3,
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the held-out data at the
        input points.
        """
        assert_params_false(self.predict_y, full_cov=full_cov, full_output_cov=full_output_cov)

        f_mean, f_var = self.predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov, cg_tolerance=cg_tolerance
        )
        return self.likelihood.predict_mean_and_var(Xnew, f_mean, f_var)

    @inherit_check_shapes
    def predict_log_density(
        self,
        data: RegressionData,
        full_cov: bool = False,
        full_output_cov: bool = False,
        cg_tolerance: Optional[float] = 1e-3,
    ) -> tf.Tensor:
        """
        Compute the log density of the data at the new data points.
        """
        assert_params_false(
            self.predict_log_density, full_cov=full_cov, full_output_cov=full_output_cov
        )

        x, y = data
        f_mean, f_var = self.predict_f(
            x, full_cov=full_cov, full_output_cov=full_output_cov, cg_tolerance=cg_tolerance
        )
        return self.likelihood.predict_log_density(x, f_mean, f_var, y)


class NystromPreconditioner:
    """
    Preconditioner of the form :math:`Q=(Q_ff + σ²I)⁻¹`,
    where L is lower triangular with :math: `LLᵀ = Kᵤᵤ`
    :math:`A = σ⁻²L⁻¹Kᵤₓ` and :math:`B = AAᵀ + I = LᵦLᵦᵀ`
    """

    @check_shapes(
        "A: [M, N]",
        "LB: [M, M]",
    )
    def __init__(self, A: tf.Tensor, LB: tf.Tensor, sigma_sq: float) -> None:
        self.A = A
        self.LB = LB
        self.sigma_sq = sigma_sq

    @check_shapes(
        "v: [B, N]",
        "return[0]: [B, N]",
        "return[1]: []",
    )
    def __call__(self, v: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Computes :math:`vᵀQ^{-1}` and `vᵀQ^{-1}v`. Note that this is
        implemented as multipication of a row vector on the right.

        :param v: Vector we want to backsolve.
        """
        sigma_sq = self.sigma_sq
        A = self.A
        LB = self.LB

        trans = tf.transpose
        trisolve = tf.linalg.triangular_solve
        matmul = tf.linalg.matmul

        v = trans(v)
        Av = matmul(A, v)
        LBinvAv = trisolve(LB, Av)
        LBinvtLBinvAv = trisolve(trans(LB), LBinvAv, lower=False)

        rv = v - matmul(A, LBinvtLBinvAv, transpose_a=True)
        vtrv = tf.reduce_sum(rv * v)
        return trans(rv) / sigma_sq, vtrv / sigma_sq


@check_shapes(
    "K: [N, N]",
    "b: [B, N]",
    "initial: [P, N]",
    "return: [P, N]",
)
def cglb_conjugate_gradient(
    K: TensorType,
    b: TensorType,
    initial: TensorType,
    preconditioner: NystromPreconditioner,
    cg_tolerance: float,
    max_steps: int,
    restart_cg_step: int,
) -> tf.Tensor:
    """
    Conjugate gradient algorithm used in CGLB model. The method of
    conjugate gradient (Hestenes and Stiefel, 1952) produces a
    sequence of vectors :math:`v_0, v_1, v_2, ..., v_N` such that
    :math:`v_0` = initial, and (in exact arithmetic)
    :math:`Kv_n = b`. In practice, the v_i often converge quickly to
    approximate :math:`K^{-1}b`, and the algorithm can be stopped
    without running N iterations.

    We assume the preconditioner, :math:`Q`, satisfies :math:`Q ≺ K`,
    and stop the algorithm when :math:`r_i = b - Kv_i` satisfies
    :math:`||rᵢᵀ||_{Q⁻¹r}^2 = rᵢᵀQ⁻¹rᵢ <= ϵ`.

    :param K: Matrix we want to backsolve from. Must be PSD.
    :param b: Vector we want to backsolve.
    :param initial: Initial vector solution.
    :param preconditioner: Preconditioner function.
    :param cg_tolerance: Expected maximum error. This value is used
        as a decision boundary against stopping criteria.
    :param max_steps: Maximum number of CG iterations.
    :param restart_cg_step: Restart step at which the CG resets the
        internal state to the initial position using the currect
        solution vector :math:`v`. Can help avoid build up of
        numerical errors.

    :return: `v` where `v` approximately satisfies :math:`Kv = b`.
    """

    class CGState(NamedTuple):
        i: tf.Tensor
        v: tf.Tensor
        r: tf.Tensor
        p: tf.Tensor
        rz: tf.Tensor

    def stopping_criterion(state: CGState) -> tf.Tensor:
        return (0.5 * state.rz > cg_tolerance) and (state.i < max_steps)

    def cg_step(state: CGState) -> List[CGState]:
        Ap = state.p @ K
        denom = tf.reduce_sum(state.p * Ap, axis=-1)
        gamma = state.rz / denom
        v = state.v + gamma * state.p
        i = state.i + 1
        r = tf.cond(
            state.i % restart_cg_step == restart_cg_step - 1,
            lambda: b - v @ K,
            lambda: state.r - gamma * Ap,
        )
        z, new_rz = preconditioner(r)
        p = tf.cond(
            state.i % restart_cg_step == restart_cg_step - 1,
            lambda: z,
            lambda: z + state.p * new_rz / state.rz,
        )
        return [CGState(i, v, r, p, new_rz)]

    Kv = initial @ K
    r = b - Kv
    z, rz = preconditioner(r)
    p = z
    i = tf.constant(0, dtype=default_int())
    initial_state = CGState(i, initial, r, p, rz)
    final_state = tf.while_loop(stopping_criterion, cg_step, [initial_state])
    final_state = tf.nest.map_structure(tf.stop_gradient, final_state)
    return final_state[0].v


# gpflow/models/gplvm.py

# Copyright 2016-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 typing import Optional

import numpy as np
import tensorflow as tf

from .. import covariances, kernels, likelihoods
from ..base import InputData, MeanAndVariance, OutputData, Parameter, RegressionData, TensorType
from ..config import default_float, default_jitter
from ..expectations import expectation
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..inducing_variables import InducingPoints
from ..kernels import Kernel
from ..mean_functions import MeanFunction, Zero
from ..probability_distributions import DiagonalGaussian
from ..utilities import assert_params_false, positive, to_default_float
from ..utilities.ops import pca_reduce
from .gpr import GPR
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import InducingVariablesLike, data_input_to_tensor, inducingpoint_wrapper


class GPLVM(GPR):
    """
    Standard GPLVM where the likelihood can be optimised with respect to the latent X.
    """

    @check_shapes(
        "data: [N, P]",
        "X_data_mean: [N, Q]",
    )
    def __init__(
        self,
        data: OutputData,
        latent_dim: int,
        X_data_mean: Optional[tf.Tensor] = None,
        kernel: Optional[Kernel] = None,
        mean_function: Optional[MeanFunction] = None,
    ):
        """
        Initialise GPLVM object. This method only works with a Gaussian likelihood.

        :param data: y data matrix.
        :param latent_dim: the number of latent dimensions (Q)
        :param X_data_mean: latent positions, for the initialisation of the latent space.
        :param kernel: kernel specification, by default Squared Exponential
        :param mean_function: mean function, by default None.
        """
        if X_data_mean is None:
            X_data_mean = pca_reduce(data, latent_dim)

        num_latent_gps = X_data_mean.shape[1]
        if num_latent_gps != latent_dim:
            msg = "Passed in number of latent {0} does not match initial X {1}."
            raise ValueError(msg.format(latent_dim, num_latent_gps))

        if mean_function is None:
            mean_function = Zero()

        if kernel is None:
            kernel = kernels.SquaredExponential(lengthscales=tf.ones((latent_dim,)))

        if data.shape[1] < num_latent_gps:
            raise ValueError("More latent dimensions than observed.")

        gpr_data = (Parameter(X_data_mean), data_input_to_tensor(data))
        super().__init__(gpr_data, kernel, mean_function=mean_function)


class BayesianGPLVM(GPModel, InternalDataTrainingLossMixin):
    @check_shapes(
        "data: [N, P]",
        "X_data_mean: [N, Q]",
        "X_data_var: [N, Q]",
        "X_prior_mean: [N, Q]",
        "X_prior_var: [N, Q]",
    )
    def __init__(
        self,
        data: OutputData,
        X_data_mean: tf.Tensor,
        X_data_var: tf.Tensor,
        kernel: Kernel,
        num_inducing_variables: Optional[int] = None,
        inducing_variable: Optional[InducingVariablesLike] = None,
        X_prior_mean: Optional[TensorType] = None,
        X_prior_var: Optional[TensorType] = None,
    ):
        """
        Initialise Bayesian GPLVM object. This method only works with a Gaussian likelihood.

        :param data: data matrix, size N (number of points) x P (dimensions)
        :param X_data_mean: initial latent positions, size N (number of points) x
            Q (latent dimensions).
        :param X_data_var: variance of latent positions ([N, Q]), for the initialisation of the
            latent space.
        :param kernel: kernel specification, by default Squared Exponential
        :param num_inducing_variables: number of inducing points, M
        :param inducing_variable: matrix of inducing points, size M (inducing points) x
            Q (latent dimensions). By default random permutation of X_data_mean.
        :param X_prior_mean: prior mean used in KL term of bound. By default 0.
            Same size as X_data_mean.
        :param X_prior_var: prior variance used in KL term of bound. By default 1.
        """
        num_data, num_latent_gps = X_data_mean.shape
        super().__init__(kernel, likelihoods.Gaussian(), num_latent_gps=num_latent_gps)
        self.data = data_input_to_tensor(data)

        self.X_data_mean = Parameter(X_data_mean)
        self.X_data_var = Parameter(X_data_var, transform=positive())

        self.num_data = num_data
        self.output_dim = self.data.shape[-1]

        if (inducing_variable is None) == (num_inducing_variables is None):
            raise ValueError(
                "BayesianGPLVM needs exactly one of `inducing_variable` and"
                " `num_inducing_variables`"
            )

        if inducing_variable is None:
            # By default we initialize by subset of initial latent points
            # Note that tf.random.shuffle returns a copy, it does not shuffle in-place
            Z = tf.random.shuffle(X_data_mean)[:num_inducing_variables]
            inducing_variable = InducingPoints(Z)

        self.inducing_variable = inducingpoint_wrapper(inducing_variable)

        assert X_data_mean.shape[1] == self.num_latent_gps

        # deal with parameters for the prior mean variance of X
        if X_prior_mean is None:
            X_prior_mean = tf.zeros((self.num_data, self.num_latent_gps), dtype=default_float())
        if X_prior_var is None:
            X_prior_var = tf.ones((self.num_data, self.num_latent_gps))

        self.X_prior_mean = tf.convert_to_tensor(np.atleast_1d(X_prior_mean), dtype=default_float())
        self.X_prior_var = tf.convert_to_tensor(np.atleast_1d(X_prior_var), dtype=default_float())

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        """
        Construct a tensorflow function to compute the bound on the marginal
        likelihood.
        """
        Y_data = self.data

        pX = DiagonalGaussian(self.X_data_mean, self.X_data_var)

        num_inducing = self.inducing_variable.num_inducing
        psi0 = tf.reduce_sum(expectation(pX, self.kernel))
        psi1 = expectation(pX, (self.kernel, self.inducing_variable))
        psi2 = tf.reduce_sum(
            expectation(
                pX, (self.kernel, self.inducing_variable), (self.kernel, self.inducing_variable)
            ),
            axis=0,
        )
        cov_uu = covariances.Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())
        L = tf.linalg.cholesky(cov_uu)
        sigma2 = self.likelihood.variance

        # Compute intermediate matrices
        A = tf.linalg.triangular_solve(L, tf.transpose(psi1), lower=True)
        tmp = tf.linalg.triangular_solve(L, psi2, lower=True)
        AAT = tf.linalg.triangular_solve(L, tf.transpose(tmp), lower=True) / sigma2
        B = AAT + tf.eye(num_inducing, dtype=default_float())
        LB = tf.linalg.cholesky(B)
        log_det_B = 2.0 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))
        c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma2

        # KL[q(x) || p(x)]
        dX_data_var = (
            self.X_data_var
            if self.X_data_var.shape.ndims == 2
            else tf.linalg.diag_part(self.X_data_var)
        )
        NQ = to_default_float(tf.size(self.X_data_mean))
        D = to_default_float(tf.shape(Y_data)[1])
        KL = -0.5 * tf.reduce_sum(tf.math.log(dX_data_var))
        KL += 0.5 * tf.reduce_sum(tf.math.log(self.X_prior_var))
        KL -= 0.5 * NQ
        KL += 0.5 * tf.reduce_sum(
            (tf.square(self.X_data_mean - self.X_prior_mean) + dX_data_var) / self.X_prior_var
        )

        # compute log marginal bound
        ND = to_default_float(tf.size(Y_data))
        bound = -0.5 * ND * tf.math.log(2 * np.pi * sigma2)
        bound += -0.5 * D * log_det_B
        bound += -0.5 * tf.reduce_sum(tf.square(Y_data)) / sigma2
        bound += 0.5 * tf.reduce_sum(tf.square(c))
        bound += -0.5 * D * (tf.reduce_sum(psi0) / sigma2 - tf.reduce_sum(tf.linalg.diag_part(AAT)))
        bound -= KL
        return bound

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points.
        Note that this is very similar to the SGPR prediction, for which
        there are notes in the SGPR notebook.

        Note: This model does not allow full output covariances.

        :param Xnew: points at which to predict
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        pX = DiagonalGaussian(self.X_data_mean, self.X_data_var)

        Y_data = self.data
        num_inducing = self.inducing_variable.num_inducing
        psi1 = expectation(pX, (self.kernel, self.inducing_variable))
        psi2 = tf.reduce_sum(
            expectation(
                pX, (self.kernel, self.inducing_variable), (self.kernel, self.inducing_variable)
            ),
            axis=0,
        )
        jitter = default_jitter()
        Kus = covariances.Kuf(self.inducing_variable, self.kernel, Xnew)
        sigma2 = self.likelihood.variance
        L = tf.linalg.cholesky(covariances.Kuu(self.inducing_variable, self.kernel, jitter=jitter))

        A = tf.linalg.triangular_solve(L, tf.transpose(psi1), lower=True)
        tmp = tf.linalg.triangular_solve(L, psi2, lower=True)
        AAT = tf.linalg.triangular_solve(L, tf.transpose(tmp), lower=True) / sigma2
        B = AAT + tf.eye(num_inducing, dtype=default_float())
        LB = tf.linalg.cholesky(B)
        c = tf.linalg.triangular_solve(LB, tf.linalg.matmul(A, Y_data), lower=True) / sigma2
        tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
        tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
        mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
        if full_cov:
            var = (
                self.kernel(Xnew)
                + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
                - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
            )
            shape = tf.stack([tf.shape(Y_data)[1], 1, 1])
            var = tf.tile(tf.expand_dims(var, 0), shape)
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), axis=0)
                - tf.reduce_sum(tf.square(tmp1), axis=0)
            )
            shape = tf.stack([1, tf.shape(Y_data)[1]])
            var = tf.tile(tf.expand_dims(var, 1), shape)
        return mean + self.mean_function(Xnew), var

    @inherit_check_shapes
    def predict_log_density(
        self, data: RegressionData, full_cov: bool = False, full_output_cov: bool = False
    ) -> tf.Tensor:
        raise NotImplementedError


# gpflow/models/gpmc.py

# Copyright 2016-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 typing import Optional

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from ..base import InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..config import default_float, default_jitter
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import assert_params_false, to_default_float
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import data_input_to_tensor


class GPMC(GPModel, InternalDataTrainingLossMixin):
    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        """
        data is a tuple of X, Y with X, a data matrix, size [N, D] and Y, a data matrix, size [N, R]
        kernel, likelihood, mean_function are appropriate GPflow objects

        This is a vanilla implementation of a GP with a non-Gaussian
        likelihood. The latent function values are represented by centered
        (whitened) variables, so

            v ~ N(0, I)
            f = Lv + m(x)

        with

            L L^T = K

        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)
        self.data = data_input_to_tensor(data)
        self.num_data = self.data[0].shape[0]
        self.V = Parameter(np.zeros((self.num_data, self.num_latent_gps)))
        self.V.prior = tfp.distributions.Normal(
            loc=to_default_float(0.0), scale=to_default_float(1.0)
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def log_posterior_density(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood() + self.log_prior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def _training_loss(self) -> tf.Tensor:  # type: ignore[override]
        return -self.log_posterior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood()

    @check_shapes(
        "return: []",
    )
    def log_likelihood(self) -> tf.Tensor:
        r"""
        Construct a tf function to compute the likelihood of a general GP
        model.

            \log p(Y | V, theta).

        """
        X_data, Y_data = self.data
        K = self.kernel(X_data)
        L = tf.linalg.cholesky(
            K + tf.eye(tf.shape(X_data)[0], dtype=default_float()) * default_jitter()
        )
        F = tf.linalg.matmul(L, self.V) + self.mean_function(X_data)

        return tf.reduce_sum(self.likelihood.log_prob(X_data, F, Y_data))

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Xnew is a data matrix, point at which we want to predict

        This method computes

            p(F* | (F=LV) )

        where F* are points on the GP at Xnew, F=LV are points on the GP at X.

        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, _Y_data = self.data
        mu, var = conditional(
            Xnew, X_data, self.kernel, self.V, full_cov=full_cov, q_sqrt=None, white=True
        )
        return mu + self.mean_function(Xnew), var


# gpflow/models/gpr.py

# Copyright 2016-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 typing import Optional

import tensorflow as tf

import gpflow

from .. import posteriors
from ..base import InputData, MeanAndVariance, RegressionData
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..likelihoods import Gaussian
from ..logdensities import multivariate_normal
from ..mean_functions import MeanFunction
from ..utilities import add_likelihood_noise_cov, assert_params_false
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import data_input_to_tensor


class GPR_deprecated(GPModel, InternalDataTrainingLossMixin):
    r"""
    Gaussian Process Regression.

    This is a vanilla implementation of GP regression with a Gaussian
    likelihood.  Multiple columns of Y are treated independently.

    The log likelihood of this model is given by

    .. math::
       \log p(Y \,|\, \mathbf f) =
            \mathcal N(Y \,|\, 0, \sigma_n^2 \mathbf{I})

    To train the model, we maximise the log _marginal_ likelihood
    w.r.t. the likelihood variance and kernel hyperparameters theta.
    The marginal likelihood is found by integrating the likelihood
    over the prior, and has the form

    .. math::
       \log p(Y \,|\, \sigma_n, \theta) =
            \mathcal N(Y \,|\, 0, \mathbf{K} + \sigma_n^2 \mathbf{I})
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        mean_function: Optional[MeanFunction] = None,
        noise_variance: Optional[float] = None,
        likelihood: Optional[Gaussian] = None,
    ):
        assert (noise_variance is None) or (
            likelihood is None
        ), "Cannot set both `noise_variance` and `likelihood`."
        if likelihood is None:
            if noise_variance is None:
                noise_variance = 1.0
            likelihood = gpflow.likelihoods.Gaussian(noise_variance)
        _, Y_data = data
        super().__init__(kernel, likelihood, mean_function, num_latent_gps=Y_data.shape[-1])
        self.data = data_input_to_tensor(data)

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_marginal_likelihood()

    @check_shapes(
        "return: []",
    )
    def log_marginal_likelihood(self) -> tf.Tensor:
        r"""
        Computes the log marginal likelihood.

        .. math::
            \log p(Y | \theta).

        """
        X, Y = self.data
        K = self.kernel(X)
        ks = add_likelihood_noise_cov(K, self.likelihood, X)
        L = tf.linalg.cholesky(ks)
        m = self.mean_function(X)

        # [R,] log-likelihoods for each independent dimension of Y
        log_prob = multivariate_normal(Y, m, L)
        return tf.reduce_sum(log_prob)

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        r"""
        This method computes predictions at X \in R^{N \x D} input points

        .. math::
            p(F* | Y)

        where F* are points on the GP at new data points, Y are noisy observations at training data
        points.
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X, Y = self.data
        err = Y - self.mean_function(X)

        kmm = self.kernel(X)
        knn = self.kernel(Xnew, full_cov=full_cov)
        kmn = self.kernel(X, Xnew)
        kmm_plus_s = add_likelihood_noise_cov(kmm, self.likelihood, X)

        conditional = gpflow.conditionals.base_conditional
        f_mean_zero, f_var = conditional(
            kmn, kmm_plus_s, knn, err, full_cov=full_cov, white=False
        )  # [N, P], [N, P] or [P, N, N]
        f_mean = f_mean_zero + self.mean_function(Xnew)
        return f_mean, f_var


class GPR_with_posterior(GPR_deprecated):
    """
    This is an implementation of GPR that provides a posterior() method that
    enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.GPRPosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """

        return posteriors.GPRPosterior(
            kernel=self.kernel,
            data=self.data,
            likelihood=self.likelihood,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, GPR's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class GPR(GPR_with_posterior):
    # subclassed to ensure __class__ == "GPR"
    pass


# gpflow/models/model.py

# Copyright 2016-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.

import abc
from typing import Any, Optional

import tensorflow as tf

from ..base import InputData, MeanAndVariance, Module, RegressionData
from ..conditionals.util import sample_mvn
from ..experimental.check_shapes import check_shapes
from ..kernels import Kernel, MultioutputKernel
from ..likelihoods import Likelihood, SwitchedLikelihood
from ..mean_functions import MeanFunction, Zero
from ..utilities import assert_params_false, to_default_float


class BayesianModel(Module, metaclass=abc.ABCMeta):
    """ Bayesian model. """

    @check_shapes(
        "return: []",
    )
    def log_prior_density(self) -> tf.Tensor:
        """
        Sum of the log prior probability densities of all (constrained) variables in this model.
        """
        if self.trainable_parameters:
            return tf.add_n([p.log_prior_density() for p in self.trainable_parameters])
        else:
            return to_default_float(0.0)

    @check_shapes(
        "return: []",
    )
    def log_posterior_density(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        """
        This may be the posterior with respect to the hyperparameters (e.g. for
        GPR) or the posterior with respect to the function (e.g. for GPMC and
        SGPMC). It assumes that maximum_log_likelihood_objective() is defined
        sensibly.
        """
        return self.maximum_log_likelihood_objective(*args, **kwargs) + self.log_prior_density()

    @check_shapes(
        "return: []",
    )
    def _training_loss(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        """
        Training loss definition. To allow MAP (maximum a-posteriori) estimation,
        adds the log density of all priors to maximum_log_likelihood_objective().
        """
        return -(self.maximum_log_likelihood_objective(*args, **kwargs) + self.log_prior_density())

    @abc.abstractmethod
    @check_shapes(
        "return: []",
    )
    def maximum_log_likelihood_objective(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        """
        Objective for maximum likelihood estimation. Should be maximized. E.g.
        log-marginal likelihood (hyperparameter likelihood) for GPR, or lower
        bound to the log-marginal likelihood (ELBO) for sparse and variational
        GPs.
        """
        raise NotImplementedError


class GPModel(BayesianModel):
    r"""
    A stateless base class for Gaussian process models, that is, those of the
    form

    .. math::
       :nowrap:

       \begin{align}
           \theta        & \sim p(\theta) \\
           f             & \sim \mathcal{GP}(m(x), k(x, x'; \theta)) \\
           f_i           & = f(x_i) \\
           y_i \,|\, f_i & \sim p(y_i|f_i)
       \end{align}

    This class mostly adds functionality for predictions. To use it, inheriting
    classes must define a predict_f function, which computes the means and
    variances of the latent function.

    These predictions are then pushed through the likelihood to obtain means
    and variances of held out data, self.predict_y.

    The predictions can also be used to compute the (log) density of held-out
    data via self.predict_log_density.

    It is also possible to draw samples from the latent GPs using
    self.predict_f_samples.
    """

    def __init__(
        self,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        super().__init__()
        assert num_latent_gps is not None, "GPModel requires specification of num_latent_gps"
        self.num_latent_gps = num_latent_gps
        if mean_function is None:
            mean_function = Zero()
        self.mean_function = mean_function
        self.kernel = kernel
        self.likelihood = likelihood

    @staticmethod
    @check_shapes(
        "data[0]: [batch..., N, D]",
        "data[1]: [batch..., N, P]",
    )
    def calc_num_latent_gps_from_data(
        data: RegressionData, kernel: Kernel, likelihood: Likelihood
    ) -> int:
        """
        Calculates the number of latent GPs required based on the data as well
        as the type of kernel and likelihood.
        """
        _, Y = data
        output_dim = Y.shape[-1]
        return GPModel.calc_num_latent_gps(kernel, likelihood, output_dim)

    @staticmethod
    def calc_num_latent_gps(kernel: Kernel, likelihood: Likelihood, output_dim: int) -> int:
        """
        Calculates the number of latent GPs required given the number of
        outputs `output_dim` and the type of likelihood and kernel.

        Note: It's not nice for `GPModel` to need to be aware of specific
        likelihoods as here. However, `num_latent_gps` is a bit more broken in
        general, we should fix this in the future. There are also some slightly
        problematic assumptions re the output dimensions of mean_function.
        See https://github.com/GPflow/GPflow/issues/1343
        """
        if isinstance(kernel, MultioutputKernel):
            # MultioutputKernels already have num_latent_gps attributes
            num_latent_gps: int = kernel.num_latent_gps
        elif isinstance(likelihood, SwitchedLikelihood):
            # the SwitchedLikelihood partitions/stitches based on the last
            # column in Y, but we should not add a separate latent GP for this!
            # hence decrement by 1
            num_latent_gps = output_dim - 1
            assert num_latent_gps > 0
        else:
            num_latent_gps = output_dim

        return num_latent_gps

    @abc.abstractmethod
    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        raise NotImplementedError

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return: [batch..., N, P] if (num_samples is None)",
        "return: [batch..., S, N, P] if (num_samples is not None)",
    )
    def predict_f_samples(
        self,
        Xnew: InputData,
        num_samples: Optional[int] = None,
        full_cov: bool = True,
        full_output_cov: bool = False,
    ) -> tf.Tensor:
        """
        Produce samples from the posterior latent function(s) at the input points.

        Currently, the method does not support `full_output_cov=True` and `full_cov=True`.

        :param Xnew:
            Input locations at which to draw samples.
        :param num_samples:
            Number of samples to draw.
            If `None`, a single sample is drawn and the return shape is [..., N, P],
            for any positive integer the return shape contains an extra batch
            dimension, [..., S, N, P], with S = num_samples and P is the number of outputs.
        :param full_cov:
            If True, draw correlated samples over the inputs. Computes the Cholesky over the
            dense covariance matrix of size [num_data, num_data].
            If False, draw samples that are uncorrelated over the inputs.
        :param full_output_cov:
            If True, draw correlated samples over the outputs.
            If False, draw samples that are uncorrelated over the outputs.
        """
        if full_cov and full_output_cov:
            raise NotImplementedError(
                "The combination of both `full_cov` and `full_output_cov` is not supported."
            )

        # check below for shape info
        mean, cov = self.predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        if full_cov:
            # mean: [..., N, P]
            # cov: [..., P, N, N]
            mean_for_sample = tf.linalg.adjoint(mean)  # [..., P, N]
            samples = sample_mvn(
                mean_for_sample, cov, full_cov, num_samples=num_samples
            )  # [..., (S), P, N]
            samples = tf.linalg.adjoint(samples)  # [..., (S), N, P]
        else:
            # mean: [..., N, P]
            # cov: [..., N, P] or [..., N, P, P]
            samples = sample_mvn(
                mean, cov, full_output_cov, num_samples=num_samples
            )  # [..., (S), N, P]
        return samples  # [..., (S), N, P]

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def predict_y(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the held-out data at the input points.
        """
        # See https://github.com/GPflow/GPflow/issues/1461
        assert_params_false(self.predict_y, full_cov=full_cov, full_output_cov=full_output_cov)

        f_mean, f_var = self.predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        return self.likelihood.predict_mean_and_var(Xnew, f_mean, f_var)

    @check_shapes(
        "data[0]: [batch..., N, D]",
        "data[1]: [batch..., N, P]",
        "return: [batch..., N]",
    )
    def predict_log_density(
        self, data: RegressionData, full_cov: bool = False, full_output_cov: bool = False
    ) -> tf.Tensor:
        """
        Compute the log density of the data at the new data points.
        """
        # See https://github.com/GPflow/GPflow/issues/1461
        assert_params_false(self.predict_y, full_cov=full_cov, full_output_cov=full_output_cov)

        X, Y = data
        f_mean, f_var = self.predict_f(X, full_cov=full_cov, full_output_cov=full_output_cov)
        return self.likelihood.predict_log_density(X, f_mean, f_var, Y)


# gpflow/models/sgpmc.py

# Copyright 2016-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 typing import Optional

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from ..base import InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import to_default_float
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import InducingPointsLike, data_input_to_tensor, inducingpoint_wrapper


class SGPMC(GPModel, InternalDataTrainingLossMixin):
    r"""
    This is the Sparse Variational GP using MCMC (SGPMC).

    The key reference is :cite:t:`hensman2015mcmc`.

    The latent function values are represented by centered
    (whitened) variables, so

    .. math::
       :nowrap:

       \begin{align}
       \mathbf v & \sim N(0, \mathbf I) \\
       \mathbf u &= \mathbf L\mathbf v
       \end{align}

    with

    .. math::
        \mathbf L \mathbf L^\top = \mathbf K


    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
        inducing_variable: Optional[InducingPointsLike] = None,
    ):
        """
        data is a tuple of X, Y with X, a data matrix, size [N, D] and Y, a data matrix, size [N, R]
        Z is a data matrix, of inducing inputs, size [M, D]
        kernel, likelihood, mean_function are appropriate GPflow objects
        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps=num_latent_gps)
        self.data = data_input_to_tensor(data)
        self.num_data = data[0].shape[0]
        self.inducing_variable = inducingpoint_wrapper(inducing_variable)
        self.V = Parameter(np.zeros((self.inducing_variable.num_inducing, self.num_latent_gps)))
        self.V.prior = tfp.distributions.Normal(
            loc=to_default_float(0.0), scale=to_default_float(1.0)
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def log_posterior_density(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood_lower_bound() + self.log_prior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def _training_loss(self) -> tf.Tensor:  # type: ignore[override]
        return -self.log_posterior_density()

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.log_likelihood_lower_bound()

    def log_likelihood_lower_bound(self) -> tf.Tensor:
        """
        This function computes the optimal density for v, q*(v), up to a constant
        """
        # get the (marginals of) q(f): exactly predicting!
        X_data, Y_data = self.data
        fmean, fvar = self.predict_f(X_data, full_cov=False)
        return tf.reduce_sum(self.likelihood.variational_expectations(X_data, fmean, fvar, Y_data))

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Xnew is a data matrix of the points at which we want to predict

        This method computes

            p(F* | (U=LV) )

        where F* are points on the GP at Xnew, F=LV are points on the GP at Z,

        """
        mu, var = conditional(
            Xnew,
            self.inducing_variable,
            self.kernel,
            self.V,
            full_cov=full_cov,
            q_sqrt=None,
            white=True,
            full_output_cov=full_output_cov,
        )
        return mu + self.mean_function(Xnew), var


# gpflow/models/sgpr.py

# Copyright 2016-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 typing import NamedTuple, Optional, Tuple

import numpy as np
import tensorflow as tf

from .. import posteriors
from ..base import InputData, MeanAndVariance, RegressionData
from ..config import default_float, default_jitter
from ..covariances.dispatch import Kuf, Kuu
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..inducing_variables import InducingPoints
from ..kernels import Kernel
from ..likelihoods import Gaussian
from ..mean_functions import MeanFunction
from ..utilities import add_noise_cov, assert_params_false, to_default_float
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import InducingPointsLike, data_input_to_tensor, inducingpoint_wrapper


class SGPRBase_deprecated(GPModel, InternalDataTrainingLossMixin):
    """
    Common base class for SGPR and GPRFITC that provides the common __init__
    and upper_bound() methods.
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        inducing_variable: InducingPointsLike,
        *,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
        noise_variance: Optional[float] = None,
        likelihood: Optional[Gaussian] = None,
    ):
        """
        This method only works with a Gaussian likelihood, its variance is
        initialized to `noise_variance`.

        :param data: a tuple of (X, Y), where the inputs X has shape [N, D]
            and the outputs Y has shape [N, R].
        :param inducing_variable:  an InducingPoints instance or a matrix of
            the pseudo inputs Z, of shape [M, D].
        :param kernel: An appropriate GPflow kernel object.
        :param mean_function: An appropriate GPflow mean function object.
        """
        assert (noise_variance is None) or (
            likelihood is None
        ), "Cannot set both `noise_variance` and `likelihood`."
        if likelihood is None:
            if noise_variance is None:
                noise_variance = 1.0
            likelihood = Gaussian(noise_variance)
        X_data, Y_data = data_input_to_tensor(data)
        num_latent_gps = Y_data.shape[-1] if num_latent_gps is None else num_latent_gps
        super().__init__(kernel, likelihood, mean_function, num_latent_gps=num_latent_gps)

        self.data = X_data, Y_data
        self.num_data = X_data.shape[0]

        self.inducing_variable: InducingPoints = inducingpoint_wrapper(inducing_variable)

    @check_shapes(
        "return: []",
    )
    def upper_bound(self) -> tf.Tensor:
        """
        Upper bound for the sparse GP regression marginal likelihood.  Note that
        the same inducing points are used for calculating the upper bound, as are
        used for computing the likelihood approximation. This may not lead to the
        best upper bound. The upper bound can be tightened by optimising Z, just
        like the lower bound. This is especially important in FITC, as FITC is
        known to produce poor inducing point locations. An optimisable upper bound
        can be found in https://github.com/markvdw/gp_upper.

        The key reference is :cite:t:`titsias_2014`.

        The key quantity, the trace term, can be computed via

        >>> _, v = conditionals.conditional(X, model.inducing_variable.Z, model.kernel,
        ...                                 np.zeros((model.inducing_variable.num_inducing, 1)))

        which computes each individual element of the trace term.
        """
        X_data, Y_data = self.data

        sigma_sq = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)  # [N]
        sigma = tf.sqrt(sigma_sq)  # [N]

        Kdiag = self.kernel(X_data, full_cov=False)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())
        kuf = Kuf(self.inducing_variable, self.kernel, X_data)

        I = tf.eye(tf.shape(kuu)[0], dtype=default_float())

        L = tf.linalg.cholesky(kuu)
        A = tf.linalg.triangular_solve(L, kuf, lower=True)

        A_sigma = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        AAT_sigma = tf.linalg.matmul(A_sigma, A_sigma, transpose_b=True)
        B = I + AAT_sigma
        LB = tf.linalg.cholesky(B)

        # Using the Trace bound, from Titsias' presentation
        c = tf.reduce_sum(Kdiag) - tf.reduce_sum(tf.square(A))

        # Alternative bound on max eigenval:
        cn_var = sigma_sq + c
        cn_std = tf.sqrt(cn_var)

        const = -0.5 * tf.reduce_sum(tf.math.log(2 * np.pi * sigma_sq))
        logdet = -tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))

        A_cn = tf.linalg.triangular_solve(L, kuf / cn_std, lower=True)
        AAT_cn = tf.linalg.matmul(A_cn, A_cn, transpose_b=True)

        err = Y_data - self.mean_function(X_data)
        LC = tf.linalg.cholesky(I + AAT_cn)
        v = tf.linalg.triangular_solve(
            LC, tf.linalg.matmul(A_cn, err / cn_std[:, None]), lower=True
        )
        quad = -0.5 * tf.reduce_sum(tf.square(err / cn_std[:, None])) + 0.5 * tf.reduce_sum(
            tf.square(v)
        )

        return const + logdet + quad


class SGPR_deprecated(SGPRBase_deprecated):
    """
    Sparse Variational GP regression.

    The key reference is :cite:t:`titsias2009variational`.
    """

    class CommonTensors(NamedTuple):
        sigma_sq: tf.Tensor
        sigma: tf.Tensor
        A: tf.Tensor
        B: tf.Tensor
        LB: tf.Tensor
        AAT: tf.Tensor
        L: tf.Tensor

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return.sigma_sq: [N]",
        "return.sigma: [N]",
        "return.A: [M, N]",
        "return.B: [M, M]",
        "return.LB: [M, M]",
        "return.AAT: [M, M]",
    )
    def _common_calculation(self) -> "SGPR.CommonTensors":
        """
        Matrices used in log-det calculation

        :return:
            * :math:`σ²`,
            * :math:`σ`,
            * :math:`A = L⁻¹K_{uf}/σ`, where :math:`LLᵀ = Kᵤᵤ`,
            * :math:`B = AAT+I`,
            * :math:`LB` where :math`LBLBᵀ = B`,
            * :math:`AAT = AAᵀ`,
        """
        x, _ = self.data  # [N]
        iv = self.inducing_variable  # [M]

        sigma_sq = tf.squeeze(self.likelihood.variance_at(x), axis=-1)  # [N]
        sigma = tf.sqrt(sigma_sq)  # [N]

        kuf = Kuf(iv, self.kernel, x)  # [M, N]
        kuu = Kuu(iv, self.kernel, jitter=default_jitter())  # [M, M]
        L = tf.linalg.cholesky(kuu)  # [M, M]

        # Compute intermediate matrices
        A = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        AAT = tf.linalg.matmul(A, A, transpose_b=True)
        B = add_noise_cov(AAT, tf.cast(1.0, AAT.dtype))
        LB = tf.linalg.cholesky(B)

        return self.CommonTensors(sigma_sq, sigma, A, B, LB, AAT, L)

    @check_shapes(
        "return: []",
    )
    def logdet_term(self, common: "SGPR.CommonTensors") -> tf.Tensor:
        r"""
        Bound from Jensen's Inequality:

        .. math::
            \log |K + σ²I| <= \log |Q + σ²I| + N * \log (1 + \textrm{tr}(K - Q)/(σ²N))

        :param common: A named tuple containing matrices that will be used
        :return: log_det, lower bound on :math:`-.5 * \textrm{output_dim} * \log |K + σ²I|`
        """
        sigma_sq = common.sigma_sq
        LB = common.LB
        AAT = common.AAT

        x, y = self.data
        outdim = to_default_float(tf.shape(y)[1])
        kdiag = self.kernel(x, full_cov=False)

        # tr(K) / σ²
        trace_k = tf.reduce_sum(kdiag / sigma_sq)
        # tr(Q) / σ²
        trace_q = tf.reduce_sum(tf.linalg.diag_part(AAT))
        # tr(K - Q) / σ²
        trace = trace_k - trace_q

        # 0.5 * log(det(B))
        half_logdet_b = tf.reduce_sum(tf.math.log(tf.linalg.diag_part(LB)))

        # sum log(σ²)
        log_sigma_sq = tf.reduce_sum(tf.math.log(sigma_sq))

        logdet_k = -outdim * (half_logdet_b + 0.5 * log_sigma_sq + 0.5 * trace)
        return logdet_k

    @check_shapes(
        "return: []",
    )
    def quad_term(self, common: "SGPR.CommonTensors") -> tf.Tensor:
        """
        :param common: A named tuple containing matrices that will be used
        :return: Lower bound on -.5 yᵀ(K + σ²I)⁻¹y
        """
        sigma = common.sigma
        A = common.A
        LB = common.LB

        x, y = self.data
        err = (y - self.mean_function(x)) / sigma[..., None]

        Aerr = tf.linalg.matmul(A, err)
        c = tf.linalg.triangular_solve(LB, Aerr, lower=True)

        # σ⁻² yᵀy
        err_inner_prod = tf.reduce_sum(tf.square(err))
        c_inner_prod = tf.reduce_sum(tf.square(c))

        quad = -0.5 * (err_inner_prod - c_inner_prod)
        return quad

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        """
        Construct a tensorflow function to compute the bound on the marginal
        likelihood. For a derivation of the terms in here, see the associated
        SGPR notebook.
        """
        common = self._common_calculation()
        output_shape = tf.shape(self.data[-1])
        num_data = to_default_float(output_shape[0])
        output_dim = to_default_float(output_shape[1])
        const = -0.5 * num_data * output_dim * np.log(2 * np.pi)
        logdet = self.logdet_term(common)
        quad = self.quad_term(common)
        return const + logdet + quad

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points
        Xnew. For a derivation of the terms in here, see the associated SGPR
        notebook.
        """
        # could copy into posterior into a fused version

        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, Y_data = self.data
        num_inducing = self.inducing_variable.num_inducing
        err = Y_data - self.mean_function(X_data)
        kuf = Kuf(self.inducing_variable, self.kernel, X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())
        Kus = Kuf(self.inducing_variable, self.kernel, Xnew)

        sigma_sq = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)
        sigma = tf.sqrt(sigma_sq)

        L = tf.linalg.cholesky(kuu)  # cache alpha, qinv
        A = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        B = tf.linalg.matmul(A, A, transpose_b=True) + tf.eye(
            num_inducing, dtype=default_float()
        )  # cache qinv
        LB = tf.linalg.cholesky(B)  # cache alpha
        Aerr = tf.linalg.matmul(A, err / sigma[..., None])
        c = tf.linalg.triangular_solve(LB, Aerr, lower=True)
        tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
        tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
        mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
        if full_cov:
            var = (
                self.kernel(Xnew)
                + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
                - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), 0)
                - tf.reduce_sum(tf.square(tmp1), 0)
            )
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        return mean + self.mean_function(Xnew), var

    @check_shapes(
        "return[0]: [M, P]",
        "return[1]: [M, M]",
    )
    def compute_qu(self) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Computes the mean and variance of q(u) = N(mu, cov), the variational distribution on
        inducing outputs.

        SVGP with this q(u) should predict identically to SGPR.

        :return: mu, cov
        """
        X_data, Y_data = self.data

        kuf = Kuf(self.inducing_variable, self.kernel, X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())

        var = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)
        std = tf.sqrt(var)
        scaled_kuf = kuf / std
        sig = kuu + tf.matmul(scaled_kuf, scaled_kuf, transpose_b=True)
        sig_sqrt = tf.linalg.cholesky(sig)

        sig_sqrt_kuu = tf.linalg.triangular_solve(sig_sqrt, kuu)

        cov = tf.linalg.matmul(sig_sqrt_kuu, sig_sqrt_kuu, transpose_a=True)
        err = Y_data - self.mean_function(X_data)
        scaled_err = err / std[..., None]
        mu = tf.linalg.matmul(
            sig_sqrt_kuu,
            tf.linalg.triangular_solve(sig_sqrt, tf.linalg.matmul(scaled_kuf, scaled_err)),
            transpose_a=True,
        )

        return mu, cov


class GPRFITC(SGPRBase_deprecated):
    """
    This implements GP regression with the FITC approximation.

    The key reference is :cite:t:`Snelson06sparsegaussian`.

    Implementation loosely based on code from GPML matlab library although
    obviously gradients are automatic in GPflow.
    """

    @check_shapes(
        "return[0]: [N, R]",
        "return[1]: [N]",
        "return[2]: [M, M]",
        "return[3]: [M, M]",
        "return[4]: [M, R]",
        "return[5]: [N, R]",
        "return[6]: [M, R]",
    )
    def common_terms(
        self,
    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
        X_data, Y_data = self.data
        num_inducing = self.inducing_variable.num_inducing
        err = Y_data - self.mean_function(X_data)  # size [N, R]
        Kdiag = self.kernel(X_data, full_cov=False)
        kuf = Kuf(self.inducing_variable, self.kernel, X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())

        sigma_sq = tf.squeeze(self.likelihood.variance_at(X_data), axis=-1)

        Luu = tf.linalg.cholesky(kuu)  # => Luu Luu^T = kuu
        V = tf.linalg.triangular_solve(Luu, kuf)  # => V^T V = Qff = kuf^T kuu^-1 kuf

        diagQff = tf.reduce_sum(tf.square(V), 0)
        nu = Kdiag - diagQff + sigma_sq

        B = tf.eye(num_inducing, dtype=default_float()) + tf.linalg.matmul(
            V / nu, V, transpose_b=True
        )
        L = tf.linalg.cholesky(B)
        beta = err / tf.expand_dims(nu, 1)  # size [N, R]
        alpha = tf.linalg.matmul(V, beta)  # size [M, R]

        gamma = tf.linalg.triangular_solve(L, alpha, lower=True)  # size [M, R]

        return err, nu, Luu, L, alpha, beta, gamma

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.fitc_log_marginal_likelihood()

    @check_shapes(
        "return: []",
    )
    def fitc_log_marginal_likelihood(self) -> tf.Tensor:
        """
        Construct a tensorflow function to compute the bound on the marginal
        likelihood.
        """

        # FITC approximation to the log marginal likelihood is
        # log ( normal( y | mean, K_fitc ) )
        # where K_fitc = Qff + diag( \nu )
        # where Qff = Kfu kuu^{-1} kuf
        # with \nu_i = Kff_{i,i} - Qff_{i,i} + \sigma^2

        # We need to compute the Mahalanobis term -0.5* err^T K_fitc^{-1} err
        # (summed over functions).

        # We need to deal with the matrix inverse term.
        # K_fitc^{-1} = ( Qff + \diag( \nu ) )^{-1}
        #             = ( V^T V + \diag( \nu ) )^{-1}
        # Applying the Woodbury identity we obtain
        #             = \diag( \nu^{-1} )
        #                 - \diag( \nu^{-1} ) V^T ( I + V \diag( \nu^{-1} ) V^T )^{-1}
        #                     V \diag(\nu^{-1} )
        # Let \beta =  \diag( \nu^{-1} ) err
        # and let \alpha = V \beta
        # then Mahalanobis term = -0.5* (
        #    \beta^T err - \alpha^T Solve( I + V \diag( \nu^{-1} ) V^T, alpha )
        # )

        err, nu, _Luu, L, _alpha, _beta, gamma = self.common_terms()

        mahalanobisTerm = -0.5 * tf.reduce_sum(
            tf.square(err) / tf.expand_dims(nu, 1)
        ) + 0.5 * tf.reduce_sum(tf.square(gamma))

        # We need to compute the log normalizing term -N/2 \log 2 pi - 0.5 \log \det( K_fitc )

        # We need to deal with the log determinant term.
        # \log \det( K_fitc ) = \log \det( Qff + \diag( \nu ) )
        #                     = \log \det( V^T V + \diag( \nu ) )
        # Applying the determinant lemma we obtain
        #                     = \log [ \det \diag( \nu ) \det( I + V \diag( \nu^{-1} ) V^T ) ]
        #                     = \log [
        #                        \det \diag( \nu ) ] + \log [ \det( I + V \diag( \nu^{-1} ) V^T )
        #                     ]

        constantTerm = -0.5 * self.num_data * tf.math.log(tf.constant(2.0 * np.pi, default_float()))
        logDeterminantTerm = -0.5 * tf.reduce_sum(tf.math.log(nu)) - tf.reduce_sum(
            tf.math.log(tf.linalg.diag_part(L))
        )
        logNormalizingTerm = constantTerm + logDeterminantTerm

        return mahalanobisTerm + logNormalizingTerm * self.num_latent_gps

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points
        Xnew.
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        _, _, Luu, L, _, _, gamma = self.common_terms()
        Kus = Kuf(self.inducing_variable, self.kernel, Xnew)  # [M, N]

        w = tf.linalg.triangular_solve(Luu, Kus, lower=True)  # [M, N]

        tmp = tf.linalg.triangular_solve(tf.transpose(L), gamma, lower=False)
        mean = tf.linalg.matmul(w, tmp, transpose_a=True) + self.mean_function(Xnew)
        intermediateA = tf.linalg.triangular_solve(L, w, lower=True)

        if full_cov:
            var = (
                self.kernel(Xnew)
                - tf.linalg.matmul(w, w, transpose_a=True)
                + tf.linalg.matmul(intermediateA, intermediateA, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                - tf.reduce_sum(tf.square(w), 0)
                + tf.reduce_sum(tf.square(intermediateA), 0)
            )  # [N, P]
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        return mean, var


class SGPR_with_posterior(SGPR_deprecated):
    """
    This is an implementation of GPR that provides a posterior() method that
    enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.SGPRPosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """

        return posteriors.SGPRPosterior(
            kernel=self.kernel,
            data=self.data,
            inducing_variable=self.inducing_variable,
            likelihood=self.likelihood,
            num_latent_gps=self.num_latent_gps,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, GPR's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class SGPR(SGPR_with_posterior):
    # subclassed to ensure __class__ == "SGPR"
    pass


# gpflow/models/svgp.py

# Copyright 2016-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 typing import Optional

import numpy as np
import tensorflow as tf

from .. import kullback_leiblers, posteriors
from ..base import AnyNDArray, InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..inducing_variables import InducingVariables
from ..kernels import Kernel
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import positive, triangular
from .model import GPModel
from .training_mixins import ExternalDataTrainingLossMixin
from .util import InducingVariablesLike, inducingpoint_wrapper


class SVGP_deprecated(GPModel, ExternalDataTrainingLossMixin):
    """
    This is the Sparse Variational GP (SVGP).

    The key reference is :cite:t:`hensman2014scalable`.
    """

    @check_shapes(
        "q_mu: [M, P]",
        "q_sqrt: [M, P] if q_diag",
        "q_sqrt: [P, M, M] if (not q_diag)",
    )
    def __init__(
        self,
        kernel: Kernel,
        likelihood: Likelihood,
        inducing_variable: InducingVariablesLike,
        *,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: int = 1,
        q_diag: bool = False,
        q_mu: Optional[tf.Tensor] = None,
        q_sqrt: Optional[tf.Tensor] = None,
        whiten: bool = True,
        num_data: Optional[tf.Tensor] = None,
    ):
        """
        - kernel, likelihood, inducing_variables, mean_function are appropriate
          GPflow objects
        - num_latent_gps is the number of latent processes to use, defaults to 1
        - q_diag is a boolean. If True, the covariance is approximated by a
          diagonal matrix.
        - whiten is a boolean. If True, we use the whitened representation of
          the inducing points.
        - num_data is the total number of observations, defaults to X.shape[0]
          (relevant when feeding in external minibatches)
        """
        # init the super class, accept args
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)
        self.num_data = num_data
        self.whiten = whiten
        self.inducing_variable: InducingVariables = inducingpoint_wrapper(inducing_variable)

        # init variational parameters
        num_inducing = self.inducing_variable.num_inducing
        self._init_variational_parameters(num_inducing, q_mu, q_sqrt, q_diag)

    @check_shapes(
        "q_mu: [M, P]",
        "q_sqrt: [M, P] if q_diag",
        "q_sqrt: [P, M, M] if (not q_diag)",
    )
    def _init_variational_parameters(
        self,
        num_inducing: int,
        q_mu: Optional[AnyNDArray],
        q_sqrt: Optional[AnyNDArray],
        q_diag: bool,
    ) -> None:
        """
        Constructs the mean and cholesky of the covariance of the variational Gaussian posterior.
        If a user passes values for `q_mu` and `q_sqrt` the routine checks if they have consistent
        and correct shapes. If a user does not specify any values for `q_mu` and `q_sqrt`, the
        routine initializes them, their shape depends on `num_inducing` and `q_diag`.

        Note: most often the comments refer to the number of observations (=output dimensions) with
        P, number of latent GPs with L, and number of inducing points M. Typically P equals L,
        but when certain multioutput kernels are used, this can change.

        Parameters
        ----------
        :param num_inducing:
            Number of inducing variables, typically refered to as M.
        :param q_mu:
            Mean of the variational Gaussian posterior. If None the function will initialise
            the mean with zeros. If not None, the shape of `q_mu` is checked.
        :param q_sqrt:
            Cholesky of the covariance of the variational Gaussian posterior.
            If None the function will initialise `q_sqrt` with identity matrix.
            If not None, the shape of `q_sqrt` is checked, depending on `q_diag`.
        :param q_diag:
            Used to check if `q_mu` and `q_sqrt` have the correct shape or to
            construct them with the correct shape. If `q_diag` is true,
            `q_sqrt` is two dimensional and only holds the square root of the
            covariance diagonal elements. If False, `q_sqrt` is three dimensional.
        """
        q_mu = np.zeros((num_inducing, self.num_latent_gps)) if q_mu is None else q_mu
        self.q_mu = Parameter(q_mu, dtype=default_float())  # [M, P]

        if q_sqrt is None:
            if q_diag:
                ones: AnyNDArray = np.ones(
                    (num_inducing, self.num_latent_gps), dtype=default_float()
                )
                self.q_sqrt = Parameter(ones, transform=positive())  # [M, P]
            else:
                np_q_sqrt: AnyNDArray = np.array(
                    [
                        np.eye(num_inducing, dtype=default_float())
                        for _ in range(self.num_latent_gps)
                    ]
                )
                self.q_sqrt = Parameter(np_q_sqrt, transform=triangular())  # [P, M, M]
        else:
            if q_diag:
                self.num_latent_gps = q_sqrt.shape[1]
                self.q_sqrt = Parameter(q_sqrt, transform=positive())  # [M, L|P]
            else:
                self.num_latent_gps = q_sqrt.shape[0]
                num_inducing = q_sqrt.shape[1]
                self.q_sqrt = Parameter(q_sqrt, transform=triangular())  # [L|P, M, M]

    @check_shapes(
        "return: []",
    )
    def prior_kl(self) -> tf.Tensor:
        return kullback_leiblers.prior_kl(
            self.inducing_variable, self.kernel, self.q_mu, self.q_sqrt, whiten=self.whiten
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self, data: RegressionData) -> tf.Tensor:  # type: ignore[override]
        return self.elbo(data)

    @check_shapes(
        "return: []",
    )
    def elbo(self, data: RegressionData) -> tf.Tensor:
        """
        This gives a variational bound (the evidence lower bound or ELBO) on
        the log marginal likelihood of the model.
        """
        X, Y = data
        kl = self.prior_kl()
        f_mean, f_var = self.predict_f(X, full_cov=False, full_output_cov=False)
        var_exp = self.likelihood.variational_expectations(X, f_mean, f_var, Y)
        if self.num_data is not None:
            num_data = tf.cast(self.num_data, kl.dtype)
            minibatch_size = tf.cast(tf.shape(X)[0], kl.dtype)
            scale = num_data / minibatch_size
        else:
            scale = tf.cast(1.0, kl.dtype)
        return tf.reduce_sum(var_exp) * scale - kl

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        mu, var = conditional(
            Xnew,
            self.inducing_variable,
            self.kernel,
            self.q_mu,
            q_sqrt=self.q_sqrt,
            full_cov=full_cov,
            white=self.whiten,
            full_output_cov=full_output_cov,
        )
        # tf.debugging.assert_positive(var)  # We really should make the tests pass with this here
        return mu + self.mean_function(Xnew), var


class SVGP_with_posterior(SVGP_deprecated):
    """
    This is the Sparse Variational GP (SVGP).

    The key reference is :cite:t:`hensman2014scalable`.

    This class provides a posterior() method that enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.BasePosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """
        return posteriors.create_posterior(
            self.kernel,
            self.inducing_variable,
            self.q_mu,
            self.q_sqrt,
            whiten=self.whiten,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, SVGP's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class SVGP(SVGP_with_posterior):
    # subclassed to ensure __class__ == "SVGP"
    pass


# gpflow/models/training_mixins.py

# Copyright 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.

"""
This module provides mixin classes to be used in conjunction with inheriting
from gpflow.models.BayesianModel (or its subclass gpflow.models.GPModel).

They provide a unified interface to obtain closures that return the training
loss, to be passed as the first argument to the minimize() method of the
optimizers defined in TensorFlow and GPflow.

All TrainingLossMixin classes assume that self._training_loss()
(which is provided by the BayesianModel base class), will be available. Note
that new models only need to implement the maximum_log_likelihood_objective
method that is defined as abstract in BayesianModel.

There are different mixins depending on whether the model already contains the
training data (InternalDataTrainingLossMixin), or requires it to be passed in
to the objective function (ExternalDataTrainingLossMixin).
"""
from typing import Callable, TypeVar, Union

import tensorflow as tf
from tensorflow.python.data.ops.iterator_ops import OwnedIterator as DatasetOwnedIterator

from ..base import InputData, OutputData, RegressionData
from ..experimental.check_shapes import check_shapes

Data = TypeVar("Data", RegressionData, InputData, OutputData)


class InternalDataTrainingLossMixin:
    """
    Mixin utility for training loss methods for models that own their own data. It provides

      - a uniform API for the training loss :meth:`training_loss`
      - a convenience method :meth:`training_loss_closure` for constructing the closure expected by
        various optimizers, namely :class:`gpflow.optimizers.Scipy` and subclasses of
        `tf.optimizers.Optimizer`.

    See :class:`ExternalDataTrainingLossMixin` for an equivalent mixin for models that do **not**
    own their own data.
    """

    @check_shapes(
        "return: []",
    )
    def training_loss(self) -> tf.Tensor:
        """
        Returns the training loss for this model.
        """
        # Type-ignore is because _training_loss should be added by implementing class.
        return self._training_loss()  # type: ignore[attr-defined]

    def training_loss_closure(self, *, compile: bool = True) -> Callable[[], tf.Tensor]:
        """
        Convenience method. Returns a closure which itself returns the training loss. This closure
        can be passed to the minimize methods on :class:`gpflow.optimizers.Scipy` and subclasses of
        `tf.optimizers.Optimizer`.

        :param compile: If `True` (default), compile the training loss function in a TensorFlow
            graph by wrapping it in tf.function()
        """
        closure = self.training_loss
        if compile:
            closure = tf.function(closure)
        return closure


class ExternalDataTrainingLossMixin:
    """
    Mixin utility for training loss methods for models that do **not** own their own data.
    It provides

      - a uniform API for the training loss :meth:`training_loss`
      - a convenience method :meth:`training_loss_closure` for constructing the closure expected by
        various optimizers, namely :class:`gpflow.optimizers.Scipy` and subclasses of
        `tf.optimizers.Optimizer`.

    See :class:`InternalDataTrainingLossMixin` for an equivalent mixin for models that **do** own
    their own data.
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
        "return: []",
    )
    def training_loss(self, data: Data) -> tf.Tensor:
        """
        Returns the training loss for this model.

        :param data: the data to be used for computing the model objective.
        """
        # Type-ignore is because _training_loss should be added by implementing class.
        return self._training_loss(data)  # type: ignore[attr-defined]

    def training_loss_closure(
        self,
        data: Union[Data, DatasetOwnedIterator],
        *,
        compile: bool = True,
    ) -> Callable[[], tf.Tensor]:
        """
        Returns a closure that computes the training loss, which by default is
        wrapped in tf.function(). This can be disabled by passing `compile=False`.

        :param data: the data to be used by the closure for computing the model
            objective. Can be the full dataset or an iterator, e.g.
            `iter(dataset.batch(batch_size))`, where dataset is an instance of
            tf.data.Dataset.
        :param compile: if True, wrap training loss in tf.function()
        """
        training_loss = self.training_loss

        if isinstance(data, DatasetOwnedIterator):
            if compile:
                # lambda because: https://github.com/GPflow/GPflow/issues/1929
                training_loss_lambda = lambda d: self.training_loss(d)
                input_signature = [data.element_spec]
                training_loss = tf.function(training_loss_lambda, input_signature=input_signature)

            def closure() -> tf.Tensor:
                assert isinstance(data, DatasetOwnedIterator)  # Hint for mypy.
                batch = next(data)
                return training_loss(batch)

        else:

            def closure() -> tf.Tensor:
                return training_loss(data)

            if compile:
                closure = tf.function(closure)

        return closure


# gpflow/models/util.py

# 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 typing import Any, Callable, Sequence, Union

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray
from ..config import default_float
from ..experimental.check_shapes import check_shapes
from ..inducing_variables import InducingPoints, InducingVariables
from .model import BayesianModel
from .training_mixins import Data, ExternalDataTrainingLossMixin

InducingVariablesLike = Union[InducingVariables, tf.Tensor, AnyNDArray]
InducingPointsLike = Union[InducingPoints, tf.Tensor, AnyNDArray]


def inducingpoint_wrapper(inducing_variable: InducingVariablesLike) -> InducingVariables:
    """
    This wrapper allows transparently passing either an InducingVariables
    object or an array specifying InducingPoints positions.
    """
    if not isinstance(inducing_variable, InducingVariables):
        inducing_variable = InducingPoints(inducing_variable)
    return inducing_variable


def _assert_equal_data(
    data1: Union[tf.Tensor, Sequence[tf.Tensor]], data2: Union[tf.Tensor, Sequence[tf.Tensor]]
) -> None:
    if isinstance(data1, tf.Tensor) and isinstance(data2, tf.Tensor):
        tf.debugging.assert_equal(data1, data2)
    else:
        for v1, v2 in zip(data1, data2):
            tf.debugging.assert_equal(v1, v2)


@check_shapes(
    "data[0]: [N, D]",
    "data[1]: [N, P]",
)
def training_loss_closure(
    model: BayesianModel, data: Data, **closure_kwargs: Any
) -> Callable[[], tf.Tensor]:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.training_loss_closure(data, **closure_kwargs)  # type: ignore[no-any-return]
    else:
        _assert_equal_data(model.data, data)
        return model.training_loss_closure(**closure_kwargs)  # type: ignore[no-any-return]


@check_shapes(
    "data[0]: [N, D]",
    "data[1]: [N, P]",
    "return: []",
)
def training_loss(model: BayesianModel, data: Data) -> tf.Tensor:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.training_loss(data)
    else:
        _assert_equal_data(model.data, data)
        return model.training_loss()


@check_shapes(
    "data[0]: [N, D]",
    "data[1]: [N, P]",
    "return: []",
)
def maximum_log_likelihood_objective(model: BayesianModel, data: Data) -> tf.Tensor:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.maximum_log_likelihood_objective(data)
    else:
        _assert_equal_data(model.data, data)
        return model.maximum_log_likelihood_objective()


def data_input_to_tensor(structure: Any) -> Any:
    """
    Converts non-tensor elements of a structure to TensorFlow tensors retaining the structure
    itself.

    The function doesn't keep original element's dtype and forcefully converts
    them to GPflow's default float type.
    """

    def convert_to_tensor(elem: Any) -> tf.Tensor:
        if tf.is_tensor(elem):
            return elem
        elif isinstance(elem, np.ndarray):
            return tf.convert_to_tensor(elem)
        return tf.convert_to_tensor(elem, dtype=default_float())

    return tf.nest.map_structure(convert_to_tensor, structure)


# gpflow/models/vgp.py

# Copyright 2016-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 typing import Optional

import numpy as np
import tensorflow as tf

import gpflow

from .. import posteriors
from ..base import InputData, MeanAndVariance, Parameter, RegressionData
from ..conditionals import conditional
from ..config import default_float, default_jitter
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from ..kernels import Kernel
from ..kullback_leiblers import gauss_kl
from ..likelihoods import Likelihood
from ..mean_functions import MeanFunction
from ..utilities import assert_params_false, is_variable, triangular, triangular_size
from .model import GPModel
from .training_mixins import InternalDataTrainingLossMixin
from .util import data_input_to_tensor


class VGP_deprecated(GPModel, InternalDataTrainingLossMixin):
    r"""
    This method approximates the Gaussian process posterior using a multivariate Gaussian.

    The idea is that the posterior over the function-value vector F is
    approximated by a Gaussian, and the KL divergence is minimised between
    the approximation and the posterior.

    This implementation is equivalent to SVGP with X=Z, but is more efficient.
    The whitened representation is used to aid optimization.

    The posterior approximation is

    .. math::

       q(\mathbf f) = N(\mathbf f \,|\, \boldsymbol \mu, \boldsymbol \Sigma)

    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        """
        data = (X, Y) contains the input points [N, D] and the observations [N, P]
        kernel, likelihood, mean_function are appropriate GPflow objects
        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)

        self.data = data_input_to_tensor(data)
        X_data, _Y_data = self.data

        static_num_data = X_data.shape[0]
        if static_num_data is None:
            q_sqrt_unconstrained_shape = (self.num_latent_gps, None)
        else:
            q_sqrt_unconstrained_shape = (self.num_latent_gps, triangular_size(static_num_data))
        self.num_data = Parameter(tf.shape(X_data)[0], shape=[], dtype=tf.int32, trainable=False)

        # Many functions below don't like `Parameter`s:
        dynamic_num_data = tf.convert_to_tensor(self.num_data)

        self.q_mu = Parameter(
            tf.zeros((dynamic_num_data, self.num_latent_gps)),
            shape=(static_num_data, num_latent_gps),
        )
        q_sqrt = tf.eye(dynamic_num_data, batch_shape=[self.num_latent_gps])
        self.q_sqrt = Parameter(
            q_sqrt,
            transform=triangular(),
            unconstrained_shape=q_sqrt_unconstrained_shape,
            constrained_shape=(num_latent_gps, static_num_data, static_num_data),
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        r"""
        This method computes the variational lower bound on the likelihood,
        which is:

            E_{q(F)} [ \log p(Y|F) ] - KL[ q(F) || p(F)]

        with

            q(\mathbf f) = N(\mathbf f \,|\, \boldsymbol \mu, \boldsymbol \Sigma)

        """
        X_data, Y_data = self.data
        num_data = tf.convert_to_tensor(self.num_data)

        # Get prior KL.
        KL = gauss_kl(self.q_mu, self.q_sqrt)

        # Get conditionals
        K = self.kernel(X_data) + tf.eye(num_data, dtype=default_float()) * default_jitter()
        L = tf.linalg.cholesky(K)
        fmean = tf.linalg.matmul(L, self.q_mu) + self.mean_function(X_data)  # [NN, ND] -> ND
        q_sqrt_dnn = tf.linalg.band_part(self.q_sqrt, -1, 0)  # [D, N, N]
        L_tiled = tf.tile(tf.expand_dims(L, 0), tf.stack([self.num_latent_gps, 1, 1]))
        LTA = tf.linalg.matmul(L_tiled, q_sqrt_dnn)  # [D, N, N]
        fvar = tf.reduce_sum(tf.square(LTA), 2)

        fvar = tf.transpose(fvar)

        # Get variational expectations.
        var_exp = self.likelihood.variational_expectations(X_data, fmean, fvar, Y_data)

        return tf.reduce_sum(var_exp) - KL

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, _Y_data = self.data
        mu, var = conditional(
            Xnew,
            X_data,
            self.kernel,
            self.q_mu,
            q_sqrt=self.q_sqrt,
            full_cov=full_cov,
            white=True,
        )
        return mu + self.mean_function(Xnew), var


class VGP_with_posterior(VGP_deprecated):
    """
    This is an implementation of VGP that provides a posterior() method that
    enables caching for faster subsequent predictions.
    """

    def posterior(
        self,
        precompute_cache: posteriors.PrecomputeCacheType = posteriors.PrecomputeCacheType.TENSOR,
    ) -> posteriors.VGPPosterior:
        """
        Create the Posterior object which contains precomputed matrices for
        faster prediction.

        precompute_cache has three settings:

        - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
          quantities and stores them as tensors (which allows differentiating
          through the prediction). This is the default.
        - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
          quantities and stores them as variables, which allows for updating
          their values without changing the compute graph (relevant for AOT
          compilation).
        - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
          immediate cache computation. This is useful for avoiding extraneous
          computations when you only want to call the posterior's
          `fused_predict_f` method.
        """
        X_data, _Y_data = self.data
        return posteriors.VGPPosterior(
            self.kernel,
            X_data,
            self.q_mu,
            self.q_sqrt,
            mean_function=self.mean_function,
            precompute_cache=precompute_cache,
        )

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        For backwards compatibility, VGP's predict_f uses the fused (no-cache)
        computation, which is more efficient during training.

        For faster (cached) prediction, predict directly from the posterior object, i.e.,:
            model.posterior().predict_f(Xnew, ...)
        """
        return self.posterior(posteriors.PrecomputeCacheType.NOCACHE).fused_predict_f(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )


class VGP(VGP_with_posterior):
    # subclassed to ensure __class__ == "VGP"
    pass


@check_shapes(
    "new_data[0]: [N, D]",
    "new_data[1]: [N, P]",
)
def update_vgp_data(vgp: VGP_deprecated, new_data: RegressionData) -> None:
    """
    Set the data on the given VGP model, and update its variational parameters.

    As opposed to many of the other models the VGP has internal parameters whose shape depends on
    the shape of the data. This functions updates the internal data of the given vgp, and updates
    the variational parameters to fit.

    This function requires that the input `vgp` were create with :class:`tf.Variable`s for
    `data`.
    """
    old_X_data, old_Y_data = vgp.data
    assert is_variable(old_X_data) and is_variable(
        old_Y_data
    ), "update_vgp_data requires the model to have been created with variable data."

    new_X_data, new_Y_data = new_data
    new_num_data = tf.shape(new_X_data)[0]
    f_mu, f_cov = vgp.predict_f(new_X_data, full_cov=True)  # [N, L], [L, N, N]

    # This model is hard-coded to use the whitened representation, i.e.  q_mu and q_sqrt
    # parametrize q(v), and u = f(X) = L v, where L = cholesky(K(X, X)) Hence we need to
    # back-transform from f_mu and f_cov to obtain the updated new_q_mu and new_q_sqrt:
    Knn = vgp.kernel(new_X_data, full_cov=True)  # [N, N]
    jitter_mat = default_jitter() * tf.eye(new_num_data, dtype=Knn.dtype)
    Lnn = tf.linalg.cholesky(Knn + jitter_mat)  # [N, N]
    new_q_mu = tf.linalg.triangular_solve(Lnn, f_mu)  # [N, L]
    tmp = tf.linalg.triangular_solve(Lnn[None], f_cov)  # [L, N, N], L⁻¹ f_cov
    S_v = tf.linalg.triangular_solve(Lnn[None], tf.linalg.matrix_transpose(tmp))  # [L, N, N]
    new_q_sqrt = tf.linalg.cholesky(S_v + jitter_mat)  # [L, N, N]

    old_X_data.assign(new_X_data)
    old_Y_data.assign(new_Y_data)
    vgp.num_data.assign(new_num_data)
    vgp.q_mu.assign(new_q_mu)
    vgp.q_sqrt.assign(new_q_sqrt)


class VGPOpperArchambeau(GPModel, InternalDataTrainingLossMixin):
    r"""
    This method approximates the Gaussian process posterior using a multivariate Gaussian.

    The key reference is :cite:t:`Opper:2009`.

    The idea is that the posterior over the function-value vector F is
    approximated by a Gaussian, and the KL divergence is minimised between
    the approximation and the posterior. It turns out that the optimal
    posterior precision shares off-diagonal elements with the prior, so
    only the diagonal elements of the precision need be adjusted.
    The posterior approximation is

    .. math::

       q(\mathbf f) = N(\mathbf f \,|\, \mathbf K \boldsymbol \alpha,
                         [\mathbf K^{-1} + \textrm{diag}(\boldsymbol \lambda))^2]^{-1})

    This approach has only 2ND parameters, rather than the N + N^2 of vgp,
    but the optimization is non-convex and in practice may cause difficulty.

    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, P]",
    )
    def __init__(
        self,
        data: RegressionData,
        kernel: Kernel,
        likelihood: Likelihood,
        mean_function: Optional[MeanFunction] = None,
        num_latent_gps: Optional[int] = None,
    ):
        """
        data = (X, Y) contains the input points [N, D] and the observations [N, P]
        kernel, likelihood, mean_function are appropriate GPflow objects
        """
        if num_latent_gps is None:
            num_latent_gps = self.calc_num_latent_gps_from_data(data, kernel, likelihood)
        super().__init__(kernel, likelihood, mean_function, num_latent_gps)

        self.data = data_input_to_tensor(data)
        X_data, _Y_data = self.data
        self.num_data = X_data.shape[0]
        self.q_alpha = Parameter(np.zeros((self.num_data, self.num_latent_gps)))
        self.q_lambda = Parameter(
            np.ones((self.num_data, self.num_latent_gps)), transform=gpflow.utilities.positive()
        )

    # type-ignore is because of changed method signature:
    @inherit_check_shapes
    def maximum_log_likelihood_objective(self) -> tf.Tensor:  # type: ignore[override]
        return self.elbo()

    @check_shapes(
        "return: []",
    )
    def elbo(self) -> tf.Tensor:
        r"""
        q_alpha, q_lambda are variational parameters, size [N, R]
        This method computes the variational lower bound on the likelihood,
        which is:

        .. math::

           E_{q(F)} [ \log p(Y|F) ] - KL[ q(F) || p(F)]

        with

        .. math::

           q(f) = N(f |
               K \alpha + \textrm{mean},
               [K^-1 + \textrm{diag}(\textrm{square}(\lambda))]^-1) .
        """
        X_data, Y_data = self.data

        K = self.kernel(X_data)
        K_alpha = tf.linalg.matmul(K, self.q_alpha)
        f_mean = K_alpha + self.mean_function(X_data)

        # compute the variance for each of the outputs
        I = tf.tile(
            tf.eye(self.num_data, dtype=default_float())[None, ...], [self.num_latent_gps, 1, 1]
        )
        A = (
            I
            + tf.transpose(self.q_lambda)[:, None, ...]
            * tf.transpose(self.q_lambda)[:, :, None, ...]
            * K
        )
        L = tf.linalg.cholesky(A)
        Li = tf.linalg.triangular_solve(L, I)
        tmp = Li / tf.transpose(self.q_lambda)[:, None, ...]
        f_var = 1.0 / tf.square(self.q_lambda) - tf.transpose(tf.reduce_sum(tf.square(tmp), 1))

        # some statistics about A are used in the KL
        A_logdet = 2.0 * tf.reduce_sum(tf.math.log(tf.linalg.diag_part(L)))
        trAi = tf.reduce_sum(tf.square(Li))

        KL = 0.5 * (
            A_logdet
            + trAi
            - self.num_data * self.num_latent_gps
            + tf.reduce_sum(K_alpha * self.q_alpha)
        )

        v_exp = self.likelihood.variational_expectations(X_data, f_mean, f_var, Y_data)
        return tf.reduce_sum(v_exp) - KL

    @inherit_check_shapes
    def predict_f(
        self, Xnew: InputData, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        r"""
        The posterior variance of F is given by

        .. math::

           q(f) = N(f |
               K \alpha + \textrm{mean}, [K^-1 + \textrm{diag}(\lambda**2)]^-1)

        Here we project this to F*, the values of the GP at Xnew which is given
        by

        .. math::

           q(F*) = N ( F* | K_{*F} \alpha + \textrm{mean}, K_{**} - K_{*f}[K_{ff} +
                                           \textrm{diag}(\lambda**-2)]^-1 K_{f*} )

        Note: This model currently does not allow full output covariances
        """
        assert_params_false(self.predict_f, full_output_cov=full_output_cov)

        X_data, _ = self.data
        # compute kernel things
        Kx = self.kernel(X_data, Xnew)
        K = self.kernel(X_data)

        # predictive mean
        f_mean = tf.linalg.matmul(Kx, self.q_alpha, transpose_a=True) + self.mean_function(Xnew)

        # predictive var
        A = K + tf.linalg.diag(tf.transpose(1.0 / tf.square(self.q_lambda)))
        L = tf.linalg.cholesky(A)
        Kx_tiled = tf.tile(Kx[None, ...], [self.num_latent_gps, 1, 1])
        LiKx = tf.linalg.triangular_solve(L, Kx_tiled)
        if full_cov:
            f_var = self.kernel(Xnew) - tf.linalg.matmul(LiKx, LiKx, transpose_a=True)
        else:
            f_var = self.kernel(Xnew, full_cov=False) - tf.reduce_sum(tf.square(LiKx), axis=1)
        return f_mean, tf.transpose(f_var)


# gpflow/monitor/__init__.py

# Copyright 2020 GPflow authors
#
# 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.
""" Provides basic functionality to monitor optimisation runs """

from .base import *
from .tensorboard import *

__all__ = [
    "ExecuteCallback",
    "ImageToTensorBoard",
    "ModelToTensorBoard",
    "Monitor",
    "MonitorTask",
    "MonitorTaskGroup",
    "ScalarToTensorBoard",
    "ToTensorBoard",
    "base",
    "tensorboard",
]


# gpflow/monitor/base.py

# Copyright 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.

""" MonitorTask base classes """

from abc import ABC, abstractmethod
from typing import Any, Callable, Collection, Union

import tensorflow as tf

__all__ = [
    "ExecuteCallback",
    "Monitor",
    "MonitorTask",
    "MonitorTaskGroup",
]


class MonitorTask(ABC):
    """
    A base class for a monitoring task.

    All monitoring tasks are callable objects.
    A descendant class must implement the `run` method, which is the body of the monitoring task.
    """

    def __call__(self, step: int, **kwargs: Any) -> None:
        """
        It calls the 'run' function and sets the current step.

        :param step: current step in the optimisation.
        :param kwargs: additional keyword arguments that can be passed
            to the `run` method of the task. This is in particular handy for
            passing keyword argument to the callback of `ScalarToTensorBoard`.
        """
        self.current_step = tf.cast(step, tf.int64)
        self.run(**kwargs)

    @abstractmethod
    def run(self, **kwargs: Any) -> None:
        """
        Implements the task to be executed on __call__.
        The current step is available through `self.current_step`.

        :param kwargs: keyword arguments available to the run method.
        """
        raise NotImplementedError


class ExecuteCallback(MonitorTask):
    """ Executes a callback as task """

    def __init__(self, callback: Callable[..., None]) -> None:
        """
        :param callback: callable to be executed during the task.
            Arguments can be passed using keyword arguments.
        """
        super().__init__()
        self.callback = callback

    def run(self, **kwargs: Any) -> None:
        self.callback(**kwargs)


class MonitorTaskGroup:
    """
    Class for grouping `MonitorTask` instances. A group defines
    all the tasks that are run at the same frequency, given by `period`.

    A `MonitorTaskGroup` can exist of a single instance or a list of
    `MonitorTask` instances.
    """

    def __init__(
        self, task_or_tasks: Union[Collection[MonitorTask], MonitorTask], period: int = 1
    ) -> None:
        """
        :param task_or_tasks: a single instance or a list of `MonitorTask` instances.
            Each `MonitorTask` in the list will be run with the given `period`.
        :param period: defines how often to run the tasks; they will execute every `period`th step.
            For large values of `period` the tasks will be less frequently run. Defaults to
            running at every step (`period = 1`).
        """
        self._tasks: Collection[MonitorTask] = []
        self.tasks = task_or_tasks  # type: ignore[assignment]
        self._period = period

    @property
    def tasks(self) -> Collection[MonitorTask]:
        return self._tasks

    @tasks.setter
    def tasks(self, task_or_tasks: Union[Collection[MonitorTask], MonitorTask]) -> None:
        """Ensures the tasks are stored as a list. Even if there is only a single task."""
        if isinstance(task_or_tasks, MonitorTask):
            self._tasks = [task_or_tasks]
        else:
            assert isinstance(task_or_tasks, Collection)
            self._tasks = list(task_or_tasks)

    def __call__(self, step: int, **kwargs: Any) -> None:
        """Call each task in the group."""
        if step % self._period == 0:
            for task in self.tasks:
                task(step, **kwargs)


class Monitor:
    r"""
    Accepts any number of of `MonitorTaskGroup` instances, and runs them
    according to their specified periodicity.

    Example use-case::

        # Create some monitor tasks
        log_dir = "logs"
        model_task = ModelToTensorBoard(log_dir, model)
        image_task = ImageToTensorBoard(log_dir, plot_prediction, "image_samples")
        lml_task = ScalarToTensorBoard(log_dir, lambda: model.log_marginal_likelihood(), "lml")

        # Plotting tasks can be quite slow, so we want to run them less frequently.
        # We group them in a `MonitorTaskGroup` and set the period to 5.
        slow_tasks = MonitorTaskGroup(image_task, period=5)

        # The other tasks are fast. We run them at each iteration of the optimisation.
        fast_tasks = MonitorTaskGroup([model_task, lml_task], period=1)

        # We pass both groups to the `Monitor`
        monitor = Monitor(fast_tasks, slow_tasks)
    """

    def __init__(self, *task_groups: MonitorTaskGroup) -> None:
        """
        :param task_groups: a list of `MonitorTaskGroup`s to be executed.
        """
        self.task_groups = task_groups

    def __call__(self, step: int, **kwargs: Any) -> None:
        for group in self.task_groups:
            group(step, **kwargs)


# gpflow/monitor/tensorboard.py

# Copyright 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.

""" Tasks that write to TensorBoard """

from io import BytesIO
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union

import numpy as np
import tensorflow as tf

from ..base import Parameter
from ..models import BayesianModel
from ..utilities import parameter_dict
from .base import MonitorTask

if TYPE_CHECKING:  # pragma: no cover
    import matplotlib


__all__ = [
    "ImageToTensorBoard",
    "ModelToTensorBoard",
    "ScalarToTensorBoard",
    "ToTensorBoard",
]


class ToTensorBoard(MonitorTask):
    writers: Dict[str, tf.summary.SummaryWriter] = {}

    def __init__(self, log_dir: str) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            Can be nested, e.g. ./logs/my_run/
        """
        super().__init__()
        if log_dir not in self.writers:
            self.writers[log_dir] = tf.summary.create_file_writer(log_dir)
        self.file_writer = self.writers[log_dir]

    def __call__(self, step: int, **kwargs: Any) -> None:
        with self.file_writer.as_default():
            super().__call__(step, **kwargs)
        self.file_writer.flush()


class ModelToTensorBoard(ToTensorBoard):
    """
    Monitoring task that creates a sensible TensorBoard for a model.

    Monitors all the model's parameters for which their name matches with `keywords_to_monitor`.
    By default, "kernel" and "likelihood" are elements of `keywords_to_monitor`.

    Example::

        keyword = "kernel", parameter = "kernel.lengthscale" => match
        keyword = "variational", parameter = "kernel.lengthscale" => no match
    """

    def __init__(
        self,
        log_dir: str,
        model: BayesianModel,
        *,
        max_size: int = 3,
        keywords_to_monitor: List[str] = ["kernel", "likelihood"],
        left_strip_character: str = ".",
    ) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            Can be a nested: for example, './logs/my_run/'.
        :param model: model to be monitord.
        :param max_size: maximum size of arrays (incl.) to store each
            element of the array independently as a scalar in the TensorBoard.
            Setting max_size to -1 will write all values. Use with care.
        :param keywords_to_monitor: specifies keywords to be monitored.
            If the parameter's name includes any of the keywords specified it
            will be monitored. By default, parameters that match the `kernel` or
            `likelihood` keyword are monitored.
            Adding a "*" to the list will match with all parameters,
            i.e. no parameters or variables will be filtered out.
        :param left_strip_character: certain frameworks prepend their variables with
            a character. GPflow adds a '.' and Keras add a '_', for example.
            When a `left_strip_character` is specified it will be stripped from the
            parameter's name. By default the '.' is left stripped, for example:
            ".likelihood.variance" becomes "likelihood.variance".
        """
        super().__init__(log_dir)
        self.model = model
        self.max_size = max_size
        self.keywords_to_monitor = keywords_to_monitor
        self.summarize_all = "*" in self.keywords_to_monitor
        self.left_strip_character = left_strip_character

    def run(self, **unused_kwargs: Any) -> None:
        for name, parameter in parameter_dict(self.model).items():
            # check if the parameter name matches any of the specified keywords
            if self.summarize_all or any(keyword in name for keyword in self.keywords_to_monitor):
                # keys are sometimes prepended with a character, which we strip
                name = name.lstrip(self.left_strip_character)
                self._summarize_parameter(name, parameter)

    def _summarize_parameter(self, name: str, param: Union[Parameter, tf.Variable]) -> None:
        """
        :param name: identifier used in tensorboard
        :param param: parameter to be stored in tensorboard
        """
        param = tf.reshape(param, (-1,))
        size = param.shape[0]

        if not isinstance(size, int):
            raise ValueError(
                f"The monitoring can not be autographed as the size of a parameter {param} "
                "is unknown at compile time. If compiling the monitor task is important, "
                "make sure the shape of all parameters is known beforehand. Otherwise, "
                "run the monitor outside the `tf.function`."
            )
        if size == 1:
            # if there's only one element do not add a numbered suffix
            tf.summary.scalar(name, param[0], step=self.current_step)
        else:
            it = range(size) if self.max_size == -1 else range(min(size, self.max_size))
            for i in it:
                tf.summary.scalar(f"{name}[{i}]", param[i], step=self.current_step)


class ScalarToTensorBoard(ToTensorBoard):
    """Stores the return value of a callback in a TensorBoard."""

    def __init__(self, log_dir: str, callback: Callable[[], float], name: str) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            For example, './logs/my_run/'.
        :param callback: callback to be executed and result written to TensorBoard.
            A callback can have arguments (e.g. data) passed to the function using
            keyword arguments.
            For example:
            ```
            lambda cb(x=None): 2 * x
            task = ScalarToTensorBoard(logdir, cb, "callback")

            # specify the argument of the function using kwargs, the names need to match.
            task(step, x=1)
            ```
        :param name: name used in TensorBoard.
        """
        super().__init__(log_dir)
        self.name = name
        self.callback = callback

    def run(self, **kwargs: Any) -> None:
        tf.summary.scalar(self.name, self.callback(**kwargs), step=self.current_step)


class ImageToTensorBoard(ToTensorBoard):
    def __init__(
        self,
        log_dir: str,
        plotting_function: Callable[
            ["matplotlib.figure.Figure", "matplotlib.figure.Axes"], "matplotlib.figure.Figure"
        ],
        name: Optional[str] = None,
        *,
        fig_kw: Optional[Dict[str, Any]] = None,
        subplots_kw: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        :param log_dir: directory in which to store the tensorboard files.
            Can be nested: for example, './logs/my_run/'.
        :param plotting_function: function performing the plotting.
        :param name: name used in TensorBoard.
        :params fig_kw: keyword arguments to be passed to Figure constructor, e.g. `figsize`.
        :params subplots_kw: keyword arguments to be passed to figure.subplots constructor, e.g.
            `nrows`, `ncols`, `sharex`, `sharey`. By default the default values
            from matplotlib.pyplot are used.
        """
        super().__init__(log_dir)
        self.plotting_function = plotting_function
        self.name = name
        self.fig_kw = fig_kw or {}
        self.subplots_kw = subplots_kw or {}

        try:
            from matplotlib.figure import Figure
        except ImportError:
            raise RuntimeError("ImageToTensorBoard requires the matplotlib package to be installed")

        self.fig = Figure(**self.fig_kw)
        if self.subplots_kw != {}:
            self.axes = self.fig.subplots(**self.subplots_kw)
        else:
            self.axes = self.fig.add_subplot(111)

    def _clear_axes(self) -> None:
        if isinstance(self.axes, np.ndarray):
            for ax in self.axes.flatten():
                ax.clear()
        else:
            self.axes.clear()

    def run(self, **unused_kwargs: Any) -> None:
        from matplotlib.backends.backend_agg import FigureCanvasAgg

        self._clear_axes()
        self.plotting_function(self.fig, self.axes)
        canvas = FigureCanvasAgg(self.fig)
        canvas.draw()

        # get PNG data from the figure
        png_buffer = BytesIO()
        canvas.print_png(png_buffer)
        png_encoded = png_buffer.getvalue()
        png_buffer.close()

        image_tensor = tf.io.decode_png(png_encoded)[None]

        # Write to TensorBoard
        tf.summary.image(self.name, image_tensor, step=self.current_step)


# gpflow/optimizers/__init__.py

from . import natgrad
from .mcmc import SamplingHelper
from .natgrad import *
from .scipy import Scipy

__all__ = [
    "NaturalGradient",
    "SamplingHelper",
    "Scipy",
    "XiNat",
    "XiSqrtMeanVar",
    "XiTransform",
    "mcmc",
    "natgrad",
    "scipy",
]


# gpflow/optimizers/mcmc.py

# Copyright 2019-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 typing import Callable, Optional, Sequence, Tuple

import tensorflow as tf

from gpflow.base import Parameter

__all__ = ["SamplingHelper"]


class SamplingHelper:
    """
    This helper makes it easy to read from variables being set with a prior and
    writes values back to the same variables.

    Example::

        model = ...  # Create a GPflow model
        hmc_helper = SamplingHelper(model.log_posterior_density, model.trainable_parameters)

        target_log_prob_fn = hmc_helper.target_log_prob_fn
        current_state = hmc_helper.current_state

        hmc = tfp.mcmc.HamiltonianMonteCarlo(target_log_prob_fn=target_log_prob_fn, ...)
        adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(hmc, ...)

        @tf.function
        def run_chain_fn():
            return mcmc.sample_chain(
                num_samples, num_burnin_steps, current_state, kernel=adaptive_hmc)

        hmc_samples = run_chain_fn()
        parameter_samples = hmc_helper.convert_to_constrained_values(hmc_samples)
    """

    def __init__(
        self, target_log_prob_fn: Callable[[], tf.Tensor], parameters: Sequence[Parameter]
    ) -> None:
        """
        :param target_log_prob_fn: a callable which returns the log-density of the model
            under the target distribution; needs to implicitly depend on the `parameters`.
            E.g. `model.log_posterior_density`.
        :param parameters: List of :class:`gpflow.Parameter` used as a state of the Markov chain.
            E.g. `model.trainable_parameters`
            Note that each parameter must have been given a prior.
        """
        if not all(isinstance(p, Parameter) and p.prior is not None for p in parameters):
            raise ValueError(
                "`parameters` should only contain gpflow.Parameter objects with priors"
            )

        self._parameters = parameters
        self._target_log_prob_fn = target_log_prob_fn
        self._variables = [p.unconstrained_variable for p in parameters]

    @property
    def current_state(self) -> Sequence[tf.Variable]:
        """Return the current state of the unconstrained variables, used in HMC."""

        return self._variables

    @property
    def target_log_prob_fn(
        self,
    ) -> Callable[..., Tuple[tf.Tensor, Callable[..., Tuple[tf.Tensor, Sequence[None]]]]]:
        """
        The target log probability, adjusted to allow for optimisation to occur on the tracked
        unconstrained underlying variables.
        """
        variables_list = self.current_state

        @tf.custom_gradient
        def _target_log_prob_fn_closure(
            *variables: tf.Variable,
        ) -> Tuple[tf.Tensor, Callable[..., Tuple[tf.Tensor, Sequence[None]]]]:
            for v_old, v_new in zip(variables_list, variables):
                v_old.assign(v_new)

            with tf.GradientTape(watch_accessed_variables=False) as tape:
                tape.watch(variables_list)
                log_prob = self._target_log_prob_fn()
                # Now need to correct for the fact that the prob fn is evaluated on the
                # constrained space while we wish to evaluate it in the unconstrained space
                for param in self._parameters:
                    if param.transform is not None:
                        x = param.unconstrained_variable
                        log_det_jacobian = param.transform.forward_log_det_jacobian(
                            x, x.shape.ndims
                        )
                        log_prob += tf.reduce_sum(log_det_jacobian)

            @tf.function
            def grad_fn(
                dy: tf.Tensor, variables: Optional[tf.Tensor] = None
            ) -> Tuple[tf.Tensor, Sequence[None]]:
                grad = tape.gradient(log_prob, variables_list)
                return grad, [None] * len(variables_list)

            return log_prob, grad_fn

        return _target_log_prob_fn_closure  # type: ignore[no-any-return]

    def convert_to_constrained_values(
        self, hmc_samples: Sequence[tf.Tensor]
    ) -> Sequence[tf.Tensor]:
        """
        Converts list of unconstrained values in `hmc_samples` to constrained
        versions. Each value in the list corresponds to an entry in parameters
        passed to the constructor; for parameters that have a transform, the
        constrained representation is returned.
        """
        values = []
        for hmc_value, param in zip(hmc_samples, self._parameters):
            if param.transform is not None:
                value = param.transform.forward(hmc_value)
            else:
                value = hmc_value
            values.append(value)
        return values


# gpflow/optimizers/natgrad.py

# Copyright 2018-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.

import abc
import functools
from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union

import numpy as np
import tensorflow as tf

from ..base import Parameter, _to_constrained
from ..experimental.check_shapes import check_shapes

Scalar = Union[float, tf.Tensor, np.ndarray]
LossClosure = Callable[[], tf.Tensor]
NatGradParameters = Union[Tuple[Parameter, Parameter], Tuple[Parameter, Parameter, "XiTransform"]]

__all__ = [
    "NaturalGradient",
    "XiNat",
    "XiSqrtMeanVar",
    "XiTransform",
]


#
# Xi transformations necessary for natural gradient optimizer.
# Abstract class and two implementations: XiNat and XiSqrtMeanVar.
#


class XiTransform(metaclass=abc.ABCMeta):
    """
    XiTransform is the base class that implements three transformations necessary
    for the natural gradient calculation wrt any parameterization.
    """

    @staticmethod
    @abc.abstractmethod
    @check_shapes(
        "mean: [N, D]",
        "varsqrt: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def meanvarsqrt_to_xi(mean: tf.Tensor, varsqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Transforms the parameter `mean` and `varsqrt` to `xi1`, `xi2`

        :param mean: the mean parameter
        :param varsqrt: the varsqrt parameter
        :return: tuple (xi1, xi2), the xi parameters
        """

    @staticmethod
    @abc.abstractmethod
    @check_shapes(
        "xi1: [N, D]",
        "xi2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def xi_to_meanvarsqrt(xi1: tf.Tensor, xi2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Transforms the parameter `xi1`, `xi2` to `mean`, `varsqrt`

        :param xi1: the ξ₁ parameter
        :param xi2: the ξ₂ parameter
        :return: tuple (mean, varsqrt), the meanvarsqrt parameters
        """

    @staticmethod
    @abc.abstractmethod
    @check_shapes(
        "nat1: [N, D]",
        "nat2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def naturals_to_xi(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        Applies the transform so that `nat1`, `nat2` is mapped to `xi1`, `xi2`

        :param nat1: the θ₁ parameter
        :param nat2: the θ₂ parameter
        :return: tuple `xi1`, `xi2`
        """


class XiNat(XiTransform):
    """
    This is the default transform. Using the natural directly saves the forward mode
    gradient, and also gives the analytic optimal solution for gamma=1 in the case
    of Gaussian likelihood.
    """

    @staticmethod
    @check_shapes(
        "mean: [N, D]",
        "varsqrt: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def meanvarsqrt_to_xi(mean: tf.Tensor, varsqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return meanvarsqrt_to_natural(mean, varsqrt)

    @staticmethod
    @check_shapes(
        "xi1: [N, D]",
        "xi2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def xi_to_meanvarsqrt(xi1: tf.Tensor, xi2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return natural_to_meanvarsqrt(xi1, xi2)

    @staticmethod
    @check_shapes(
        "nat1: [N, D]",
        "nat2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def naturals_to_xi(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return nat1, nat2


class XiSqrtMeanVar(XiTransform):
    """
    This transformation will perform natural gradient descent on the model parameters,
    so saves the conversion to and from Xi.
    """

    @staticmethod
    @check_shapes(
        "mean: [N, D]",
        "varsqrt: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def meanvarsqrt_to_xi(mean: tf.Tensor, varsqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return mean, varsqrt

    @staticmethod
    @check_shapes(
        "xi1: [N, D]",
        "xi2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def xi_to_meanvarsqrt(xi1: tf.Tensor, xi2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return xi1, xi2

    @staticmethod
    @check_shapes(
        "nat1: [N, D]",
        "nat2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def naturals_to_xi(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
        return natural_to_meanvarsqrt(nat1, nat2)


class NaturalGradient(tf.optimizers.Optimizer):
    """
    Implements a natural gradient descent optimizer for variational models
    that are based on a distribution q(u) = N(q_mu, q_sqrt q_sqrtᵀ) that is
    parameterized by mean q_mu and lower-triangular Cholesky factor q_sqrt
    of the covariance.

    Note that this optimizer does not implement the standard API of
    tf.optimizers.Optimizer. Its only public method is minimize(), which has
    a custom signature (var_list needs to be a list of (q_mu, q_sqrt) tuples,
    where q_mu and q_sqrt are gpflow.Parameter instances, not tf.Variable).

    Note furthermore that the natural gradients are implemented only for the
    full covariance case (i.e., q_diag=True is NOT supported).

    When using in your work, please cite :cite:t:`salimbeni18`.
    """

    def __init__(
        self, gamma: Scalar, xi_transform: XiTransform = XiNat(), name: Optional[str] = None
    ) -> None:
        """
        :param gamma: natgrad step length
        :param xi_transform: default ξ transform (can be overridden in the call to minimize())
            The XiNat default choice works well in general.
        """
        name = self.__class__.__name__ if name is None else name
        super().__init__(name)
        self.gamma = gamma
        self.xi_transform = xi_transform

    @check_shapes(
        "var_list[all][0]: [N, D]",
        "var_list[all][1]: [D, N, N]",
    )
    def minimize(
        self,
        loss_fn: LossClosure,
        var_list: Sequence[NatGradParameters],
    ) -> None:
        """
        Minimizes objective function of the model.
        Natural Gradient optimizer works with variational parameters only.

        GPflow implements the `XiNat` (default) and `XiSqrtMeanVar` transformations
        for parameters. Custom transformations that implement the `XiTransform`
        interface are also possible.

        :param loss_fn: Loss function.
        :param var_list: List of pair tuples of variational parameters or
            triplet tuple with variational parameters and ξ transformation.
            If ξ is not specified, will use self.xi_transform.
            For example, `var_list` could be::

                var_list = [
                    (q_mu1, q_sqrt1),
                    (q_mu2, q_sqrt2, XiSqrtMeanVar())
                ]
        """
        parameters = [(v[0], v[1], (v[2] if len(v) > 2 else None)) for v in var_list]  # type: ignore[misc]
        self._natgrad_steps(loss_fn, parameters)

    @check_shapes(
        "parameters[all][0]: [N, D]",
        "parameters[all][1]: [D, N, N]",
    )
    def _natgrad_steps(
        self,
        loss_fn: LossClosure,
        parameters: Sequence[Tuple[Parameter, Parameter, Optional[XiTransform]]],
    ) -> None:
        """
        Computes gradients of loss_fn() w.r.t. q_mu and q_sqrt, and updates
        these parameters using the natgrad backwards step, for all sets of
        variational parameters passed in.

        :param loss_fn: Loss function.
        :param parameters: List of tuples (q_mu, q_sqrt, xi_transform)
        """
        q_mus, q_sqrts, xis = zip(*parameters)
        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(q_mu_vars + q_sqrt_vars)
            loss = loss_fn()

        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"):
            for q_mu_grad, q_sqrt_grad, q_mu, q_sqrt, xi_transform in zip(
                q_mu_grads, q_sqrt_grads, q_mus, q_sqrts, xis
            ):
                self._natgrad_apply_gradients(q_mu_grad, q_sqrt_grad, q_mu, q_sqrt, xi_transform)

    @check_shapes(
        "q_mu_grad: [N, D]",
        "q_sqrt_grad: [D, N_N_transformed...]",
        "q_mu: [N, D]",
        "q_sqrt: [D, N, N]",
    )
    def _natgrad_apply_gradients(
        self,
        q_mu_grad: tf.Tensor,
        q_sqrt_grad: tf.Tensor,
        q_mu: Parameter,
        q_sqrt: Parameter,
        xi_transform: Optional[XiTransform] = None,
    ) -> None:
        """
        This function does the backward step on the q_mu and q_sqrt parameters,
        given the gradients of the loss function with respect to their unconstrained
        variables. I.e., it expects the arguments to come from

            with tf.GradientTape() as tape:
                loss = loss_function()
            q_mu_grad, q_mu_sqrt = tape.gradient(loss, [q_mu, q_sqrt])

        (Note that tape.gradient() returns the gradients in *unconstrained* space!)

        Implements equation [10] from :cite:t:`salimbeni18`.

        In addition, for convenience with the rest of GPflow, this code computes ∂L/∂η using
        the chain rule (the following assumes a numerator layout where the gradient is a row
        vector; note that TensorFlow actually returns a column vector), where L is the loss:

        ∂L/∂η = (∂L / ∂[q_mu, q_sqrt])(∂[q_mu, q_sqrt] / ∂η)

        In total there are three derivative calculations:
        natgrad of L w.r.t ξ  = (∂ξ / ∂θ) [(∂L / ∂[q_mu, q_sqrt]) (∂[q_mu, q_sqrt] / ∂η)]ᵀ

        Note that if ξ = θ (i.e. [q_mu, q_sqrt]) some of these calculations are the identity.
        In the code η = eta, ξ = xi, θ = nat.

        :param q_mu_grad: gradient of loss w.r.t. q_mu (in unconstrained space)
        :param q_sqrt_grad: gradient of loss w.r.t. q_sqrt (in unconstrained space)
        :param q_mu: parameter for the mean of q(u) with shape [M, L]
        :param q_sqrt: parameter for the square root of the covariance of q(u)
            with shape [L, M, M] (the diagonal parametrization, q_diag=True, is NOT supported)
        :param xi_transform: the ξ transform to use (self.xi_transform if not specified)
        """
        if xi_transform is None:
            xi_transform = self.xi_transform

        # 1) the ordinary gpflow gradient
        dL_dmean = _to_constrained(q_mu_grad, q_mu.transform)
        dL_dvarsqrt = _to_constrained(q_sqrt_grad, q_sqrt.transform)

        with tf.GradientTape(persistent=True, watch_accessed_variables=False) as tape:
            tape.watch([q_mu.unconstrained_variable, q_sqrt.unconstrained_variable])

            # the three parameterizations as functions of [q_mu, q_sqrt]
            eta1, eta2 = meanvarsqrt_to_expectation(q_mu, q_sqrt)
            # we need these to calculate the relevant gradients
            meanvarsqrt = expectation_to_meanvarsqrt(eta1, eta2)

            if not isinstance(xi_transform, XiNat):
                nat1, nat2 = meanvarsqrt_to_natural(q_mu, q_sqrt)
                xi1_nat, xi2_nat = xi_transform.naturals_to_xi(nat1, nat2)
                dummy_tensors = tf.ones_like(xi1_nat), tf.ones_like(xi2_nat)
                with tf.GradientTape(watch_accessed_variables=False) as forward_tape:
                    forward_tape.watch(dummy_tensors)
                    dummy_gradients = tape.gradient(
                        [xi1_nat, xi2_nat], [nat1, nat2], output_gradients=dummy_tensors
                    )

        # 2) the chain rule to get ∂L/∂η, where η (eta) are the expectation parameters
        dL_deta1, dL_deta2 = tape.gradient(
            meanvarsqrt, [eta1, eta2], output_gradients=[dL_dmean, dL_dvarsqrt]
        )

        if not isinstance(xi_transform, XiNat):
            nat_dL_xi1, nat_dL_xi2 = forward_tape.gradient(
                dummy_gradients, dummy_tensors, output_gradients=[dL_deta1, dL_deta2]
            )
        else:
            nat_dL_xi1, nat_dL_xi2 = dL_deta1, dL_deta2

        del tape  # Remove "persistent" tape

        xi1, xi2 = xi_transform.meanvarsqrt_to_xi(q_mu, q_sqrt)
        xi1_new = xi1 - self.gamma * nat_dL_xi1
        xi2_new = xi2 - self.gamma * nat_dL_xi2

        # Transform back to the model parameters [q_mu, q_sqrt]
        mean_new, varsqrt_new = xi_transform.xi_to_meanvarsqrt(xi1_new, xi2_new)

        q_mu.assign(mean_new)
        q_sqrt.assign(varsqrt_new)

    def get_config(self) -> Dict[str, Any]:
        config: Dict[str, Any] = super().get_config()
        config.update({"gamma": self._serialize_hyperparameter("gamma")})
        return config


#
# Auxiliary gaussian parameter conversion functions.
#
# The following functions expect their first and second inputs to have shape
# [D, N, 1] and [D, N, N], respectively. Return values are also of shapes [D, N, 1] and [D, N, N].


def swap_dimensions(
    method: Callable[[tf.Tensor, tf.Tensor], Tuple[tf.Tensor, tf.Tensor]]
) -> Callable[..., Tuple[tf.Tensor, tf.Tensor]]:
    """
    Converts between GPflow indexing and tensorflow indexing
    `method` is a function that broadcasts over the first dimension (i.e. like all tensorflow matrix
    ops):

    * `method` inputs [D, N, 1], [D, N, N]
    * `method` outputs [D, N, 1], [D, N, N]

    :return: Function that broadcasts over the final dimension (i.e. compatible with GPflow):

        * inputs: [N, D], [D, N, N]
        * outputs: [N, D], [D, N, N]
    """

    @functools.wraps(method)
    @check_shapes(
        "a_nd: [N, D] if swap",
        "a_nd: [D, N, 1] if not swap",
        "b_dnn: [D, N, N]",
        "return[0]: [N, D] if swap",
        "return[0]: [D, N, 1] if not swap",
        "return[1]: [D, N, N]",
    )
    def wrapper(
        a_nd: tf.Tensor, b_dnn: tf.Tensor, swap: bool = True
    ) -> Tuple[tf.Tensor, tf.Tensor]:
        if swap:
            a_dn1 = tf.linalg.adjoint(a_nd)[:, :, None]
            A_dn1, B_dnn = method(a_dn1, b_dnn)
            A_nd = tf.linalg.adjoint(A_dn1[:, :, 0])
            return A_nd, B_dnn
        else:
            return method(a_nd, b_dnn)

    return wrapper


@swap_dimensions
@check_shapes(
    "nat1: [D, N, 1]",
    "nat2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def natural_to_meanvarsqrt(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    var_sqrt_inv = tf.linalg.cholesky(-2 * nat2)
    var_sqrt = _inverse_lower_triangular(var_sqrt_inv)
    S = tf.linalg.matmul(var_sqrt, var_sqrt, transpose_a=True)
    mu = tf.linalg.matmul(S, nat1)
    # We need the decomposition of S as L L^T, not as L^T L,
    # hence we need another cholesky.
    return mu, tf.linalg.cholesky(S)


@swap_dimensions
@check_shapes(
    "mu: [D, N, 1]",
    "s_sqrt: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def meanvarsqrt_to_natural(mu: tf.Tensor, s_sqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    s_sqrt_inv = _inverse_lower_triangular(s_sqrt)
    s_inv = tf.linalg.matmul(s_sqrt_inv, s_sqrt_inv, transpose_a=True)
    return tf.linalg.matmul(s_inv, mu), -0.5 * s_inv


@swap_dimensions
@check_shapes(
    "nat1: [D, N, 1]",
    "nat2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def natural_to_expectation(nat1: tf.Tensor, nat2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    args = natural_to_meanvarsqrt(nat1, nat2, swap=False)
    return meanvarsqrt_to_expectation(*args, swap=False)


@swap_dimensions
@check_shapes(
    "eta1: [D, N, 1]",
    "eta2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def expectation_to_natural(eta1: tf.Tensor, eta2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    args = expectation_to_meanvarsqrt(eta1, eta2, swap=False)
    return meanvarsqrt_to_natural(*args, swap=False)


@swap_dimensions
@check_shapes(
    "eta1: [D, N, 1]",
    "eta2: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def expectation_to_meanvarsqrt(eta1: tf.Tensor, eta2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    var = eta2 - tf.linalg.matmul(eta1, eta1, transpose_b=True)
    return eta1, tf.linalg.cholesky(var)


@swap_dimensions
@check_shapes(
    "m: [D, N, 1]",
    "v_sqrt: [D, N, N]",
    "return[0]: [D, N, 1]",
    "return[1]: [D, N, N]",
)
def meanvarsqrt_to_expectation(m: tf.Tensor, v_sqrt: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
    v = tf.linalg.matmul(v_sqrt, v_sqrt, transpose_b=True)
    return m, v + tf.linalg.matmul(m, m, transpose_b=True)


@check_shapes(
    "M: [D, N, N]",
    "return: [D, N, N]",
)
def _inverse_lower_triangular(M: tf.Tensor) -> tf.Tensor:
    """
    Take inverse of lower triangular (e.g. Cholesky) matrix. This function
    broadcasts over the first index.

    :param M: Tensor with lower triangular structure of shape [D, N, N]
    :return: The inverse of the Cholesky decomposition. Same shape as input.
    """
    if M.shape.ndims != 3:  # pragma: no cover
        raise ValueError("Number of dimensions for input is required to be 3.")
    D, N = tf.shape(M)[0], tf.shape(M)[1]
    I_dnn = tf.eye(N, dtype=M.dtype)[None, :, :] * tf.ones((D, 1, 1), dtype=M.dtype)
    return tf.linalg.triangular_solve(M, I_dnn)


# gpflow/optimizers/scipy.py

# 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.

import warnings
from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union

import numpy as np
import scipy.optimize
import tensorflow as tf
from scipy.optimize import OptimizeResult

from ..base import AnyNDArray
from ..monitor.base import Monitor

__all__ = ["Scipy"]

Variables = Iterable[tf.Variable]  # deprecated
StepCallback = Union[Callable[[int, Sequence[tf.Variable], Sequence[tf.Tensor]], None], Monitor]
LossClosure = Callable[[], tf.Tensor]


class Scipy:
    def minimize(
        self,
        closure: LossClosure,
        variables: Sequence[tf.Variable],
        method: Optional[str] = "L-BFGS-B",
        step_callback: Optional[StepCallback] = None,
        compile: bool = True,
        allow_unused_variables: bool = False,
        **scipy_kwargs: Any,
    ) -> OptimizeResult:
        """
        Minimize `closure`.

        Minimize is a wrapper around the `scipy.optimize.minimize` function handling the packing and
        unpacking of a list of shaped variables on the TensorFlow side vs. the flat numpy array
        required on the Scipy side.

        :param closure: A closure that re-evaluates the model, returning the loss to be minimized.
        :param variables: The list (tuple) of variables to be optimized
            (typically `model.trainable_variables`)
        :param method: The type of solver to use in SciPy. Defaults to "L-BFGS-B".
        :param step_callback: If not None, a callable that gets called once after each optimisation
            step. The callable is passed the arguments `step`, `variables`, and `values`. `step` is
            the optimisation step counter, `variables` is the list of trainable variables as above,
            and `values` is the corresponding list of tensors of matching shape that contains their
            value at this optimisation step.
        :param compile: If True, wraps the evaluation function (the passed `closure` as well as its
            gradient computation) inside a `tf.function()`, which will improve optimization speed in
            most cases.
        :param allow_unused_variables: Whether to allow variables that are not actually used in the
            closure.
        :param scipy_kwargs: Arguments passed through to `scipy.optimize.minimize`.
            Note that Scipy's minimize() takes a `callback` argument, but you probably want to use
            our wrapper and pass in `step_callback`.
        :returns:
            The optimization result represented as a Scipy ``OptimizeResult`` object.
            See the Scipy documentation for description of attributes.
        """
        if not callable(closure):
            raise TypeError(
                "The 'closure' argument is expected to be a callable object."
            )  # pragma: no cover
        variables = tuple(variables)
        if not all(isinstance(v, tf.Variable) for v in variables):
            raise TypeError(
                "The 'variables' argument is expected to only contain tf.Variable instances"
                " (use model.trainable_variables, not model.trainable_parameters)"
            )  # pragma: no cover
        initial_params = self.initial_parameters(variables)

        func = self.eval_func(
            closure, variables, compile=compile, allow_unused_variables=allow_unused_variables
        )
        if step_callback is not None:
            if "callback" in scipy_kwargs:
                raise ValueError("Callback passed both via `step_callback` and `callback`")

            callback = self.callback_func(variables, step_callback)
            scipy_kwargs.update(dict(callback=callback))

        return scipy.optimize.minimize(
            func, initial_params, jac=True, method=method, **scipy_kwargs
        )

    @classmethod
    def initial_parameters(cls, variables: Sequence[tf.Variable]) -> tf.Tensor:
        return cls.pack_tensors(variables)

    @classmethod
    def eval_func(
        cls,
        closure: LossClosure,
        variables: Sequence[tf.Variable],
        compile: bool = True,
        allow_unused_variables: bool = False,
    ) -> Callable[[AnyNDArray], Tuple[AnyNDArray, AnyNDArray]]:
        first_call = True

        def _tf_eval(x: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
            nonlocal first_call

            values = cls.unpack_tensors(variables, x)
            cls.assign_tensors(variables, values)

            if first_call:
                # Only check for unconnected gradients on the first function evaluation.
                loss, grads = _compute_loss_and_gradients(
                    closure, variables, tf.UnconnectedGradients.NONE
                )
                grads = cls._filter_unused_variables(variables, grads, allow_unused_variables)
                first_call = False
            else:
                loss, grads = _compute_loss_and_gradients(
                    closure, variables, tf.UnconnectedGradients.ZERO
                )

            return loss, cls.pack_tensors(grads)

        if compile:
            _tf_eval = tf.function(_tf_eval)

        def _eval(x: AnyNDArray) -> Tuple[AnyNDArray, AnyNDArray]:
            loss, grad = _tf_eval(tf.convert_to_tensor(x))
            return loss.numpy().astype(np.float64), grad.numpy().astype(np.float64)

        return _eval

    @staticmethod
    def _filter_unused_variables(
        variables: Sequence[tf.Variable], grads: Sequence[tf.Tensor], allow_unused_variables: bool
    ) -> Sequence[tf.Tensor]:
        filtered_grads = []
        unused_variables = []
        for i, grad in enumerate(grads):
            if grad is None:
                variable = variables[i]
                filtered_grads.append(tf.zeros_like(variable))
                unused_variables.append(variable.name)
            else:
                filtered_grads.append(grad)

        if unused_variables:
            msg = (
                "Some variables does not have a gradient, and appear unused in / not connected to"
                f" the loss closure: {unused_variables}."
            )
            if allow_unused_variables:
                warnings.warn(msg)
            else:
                raise ValueError(msg)

        return filtered_grads

    @classmethod
    def callback_func(
        cls, variables: Sequence[tf.Variable], step_callback: StepCallback
    ) -> Callable[[AnyNDArray], None]:
        step = 0  # type: int

        def _callback(x: AnyNDArray) -> None:
            nonlocal step

            if isinstance(step_callback, Monitor):
                step_callback(step)
            else:
                values = cls.unpack_tensors(variables, x)
                step_callback(step, variables, values)

            step += 1

        return _callback

    @staticmethod
    def pack_tensors(tensors: Sequence[Union[tf.Tensor, tf.Variable]]) -> tf.Tensor:
        flats = [tf.reshape(tensor, (-1,)) for tensor in tensors]
        tensors_vector = tf.concat(flats, axis=0)
        return tensors_vector

    @staticmethod
    def unpack_tensors(
        to_tensors: Sequence[Union[tf.Tensor, tf.Variable]], from_vector: tf.Tensor
    ) -> List[tf.Tensor]:
        s = 0
        values = []
        for target_tensor in to_tensors:
            shape = tf.shape(target_tensor)
            dtype = target_tensor.dtype
            tensor_size = tf.reduce_prod(shape)
            tensor_vector = from_vector[s : s + tensor_size]
            tensor = tf.reshape(tf.cast(tensor_vector, dtype), shape)
            values.append(tensor)
            s += tensor_size
        return values

    @staticmethod
    def assign_tensors(to_tensors: Sequence[tf.Variable], values: Sequence[tf.Tensor]) -> None:
        if len(to_tensors) != len(values):
            raise ValueError("to_tensors and values should have same length")
        for target, value in zip(to_tensors, values):
            target.assign(value)


def _compute_loss_and_gradients(
    loss_closure: LossClosure,
    variables: Sequence[tf.Variable],
    unconnected_gradients: tf.UnconnectedGradients,
) -> Tuple[tf.Tensor, Sequence[tf.Tensor]]:
    with tf.GradientTape(watch_accessed_variables=False) as tape:
        tape.watch(variables)
        loss = loss_closure()
    grads = tape.gradient(loss, variables, unconnected_gradients=unconnected_gradients)
    return loss, grads


# gpflow/quadrature/__init__.py

from .base import GaussianQuadrature
from .deprecated import hermgauss, mvhermgauss, mvnquad, ndiag_mc, ndiagquad
from .gauss_hermite import NDiagGHQuadrature

__all__ = [
    "GaussianQuadrature",
    "NDiagGHQuadrature",
    "base",
    "deprecated",
    "gauss_hermite",
    "hermgauss",
    "mvhermgauss",
    "mvnquad",
    "ndiag_mc",
    "ndiagquad",
]


# gpflow/quadrature/base.py

# Copyright 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.

import abc
from typing import Any, Callable, Iterable, Tuple, Union

import tensorflow as tf

from ..base import TensorType
from ..experimental.check_shapes import check_shapes


class GaussianQuadrature:
    """
    Abstract class implementing quadrature methods to compute Gaussian Expectations.
    Inheriting classes must provide the method _build_X_W to create points and weights
    to be used for quadrature.
    """

    @abc.abstractmethod
    @check_shapes(
        "mean: [batch..., D]",
        "var: [batch..., D]",
        "return[0]: [N_quad_points, batch..., D]",
        "return[1]: [N_quad_points, broadcast batch..., 1]",
    )
    def _build_X_W(self, mean: TensorType, var: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        raise NotImplementedError

    @check_shapes(
        "mean: [batch..., D]",
        "var: [batch..., D]",
        "return: [n_funs..., batch..., broadcast D]",
    )
    def __call__(
        self,
        fun: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
        mean: TensorType,
        var: TensorType,
        *args: Any,
        **kwargs: Any,
    ) -> tf.Tensor:
        r"""
        Compute the Gaussian Expectation of a function f::

            X ~ N(mean, var)
            E[f(X)] = ∫f(x, *args, **kwargs)p(x)dx

        Using the formula::

            E[f(X)] = sum_{i=1}^{N_quad_points} f(x_i) * w_i

        where x_i, w_i must be provided by the inheriting class through self._build_X_W.

        :param fun: Callable or Iterable of Callables that operates elementwise, with
            signature f(X, \*args, \*\*kwargs). Moreover, it must satisfy the shape-mapping::

                X shape: [N_quad_points, batch..., d].
                f(X) shape: [N_quad_points, batch..., broadcast d].

            In most cases, f should only operate over the last dimension of X
        :param mean: Array/Tensor representing the mean of a d-Variate Gaussian distribution
        :param var: Array/Tensor representing the variance of a d-Variate Gaussian distribution
        :param args: Passed to fun
        :param kargs: Passed to fun
        :return: Gaussian expectation of fun
        """

        X, W = self._build_X_W(mean, var)
        if isinstance(fun, Iterable):
            return [tf.reduce_sum(f(X, *args, **kwargs) * W, axis=0) for f in fun]
        return tf.reduce_sum(fun(X, *args, **kwargs) * W, axis=0)

    @check_shapes(
        "mean: [batch..., D]",
        "var: [batch..., D]",
        "return: [batch..., broadcast D]",
    )
    def logspace(
        self,
        fun: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
        mean: TensorType,
        var: TensorType,
        *args: Any,
        **kwargs: Any,
    ) -> tf.Tensor:
        r"""
        Compute the Gaussian log-Expectation of a the exponential of a function f::

            X ~ N(mean, var)
            log E[exp[f(X)]] = log ∫exp[f(x, *args, **kwargs)]p(x)dx

        Using the formula::

            log E[exp[f(X)]] = log sum_{i=1}^{N_quad_points} exp[f(x_i) + log w_i]

        where x_i, w_i must be provided by the inheriting class through self._build_X_W.
        The computations broadcast along batch-dimensions, represented by [batch...].

        :param fun: Callable or Iterable of Callables that operates elementwise, with
            signature f(X, \*args, \*\*kwargs). Moreover, it must satisfy the shape-mapping::

                X shape: [N_quad_points, batch..., d].
                f(X) shape: [N_quad_points, batch..., broadcast d].

            In most cases, f should only operate over the last dimension of X
        :param mean: Array/Tensor representing the mean of a d-Variate Gaussian distribution
        :param var: Array/Tensor representing the variance of a d-Variate Gaussian distribution
        :param args: Passed to fun
        :param kwargs: Passed to fun
        :return: Gaussian log-expectation of the exponential of a function f
        """

        X, W = self._build_X_W(mean, var)
        logW = tf.math.log(W)
        if isinstance(fun, Iterable):
            return [tf.reduce_logsumexp(f(X, *args, **kwargs) + logW, axis=0) for f in fun]
        return tf.reduce_logsumexp(fun(X, *args, **kwargs) + logW, axis=0)


# gpflow/quadrature/deprecated.py

# 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.

import itertools
import warnings
from functools import wraps
from typing import Callable, Iterable, List, Optional, Tuple, Union

import numpy as np
import tensorflow as tf

from ..base import AnyNDArray, TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes
from ..utilities import to_default_float
from .gauss_hermite import NDiagGHQuadrature


@check_shapes(
    "return[0]: [n_quad_points]",
    "return[1]: [n_quad_points]",
)
def hermgauss(n: int) -> Tuple[AnyNDArray, AnyNDArray]:
    # Type-ignore is because for some versions mypy can't find np.polynomial.hermite
    x, w = np.polynomial.hermite.hermgauss(n)
    x, w = x.astype(default_float()), w.astype(default_float())
    return x, w


@check_shapes(
    "return[0]: [n_quad_points, D]",
    "return[1]: [n_quad_points]",
)
def mvhermgauss(H: int, D: int) -> Tuple[AnyNDArray, AnyNDArray]:
    """
    Return the evaluation locations 'xn', and weights 'wn' for a multivariate
    Gauss-Hermite quadrature.

    The outputs can be used to approximate the following type of integral:
    int exp(-x)*f(x) dx ~ sum_i w[i,:]*f(x[i,:])

    :param H: Number of Gauss-Hermite evaluation points.
    :param D: Number of input dimensions. Needs to be known at call-time.
    :return: eval_locations 'x' (H**DxD), weights 'w' (H**D)
    """
    gh_x, gh_w = hermgauss(H)
    x: AnyNDArray = np.array(list(itertools.product(*(gh_x,) * D)))  # H**DxD
    w = np.prod(np.array(list(itertools.product(*(gh_w,) * D))), 1)  # H**D
    return x, w


@check_shapes(
    "means: [N, Din]",
    "covs: [N, Din, Din]",
    "return: [N, Dout...]",
)
def mvnquad(
    func: Callable[[tf.Tensor], tf.Tensor],
    means: TensorType,
    covs: TensorType,
    H: int,
    Din: Optional[int] = None,
    Dout: Optional[Tuple[int, ...]] = None,
) -> tf.Tensor:
    """
    Computes N Gaussian expectation integrals of a single function 'f'
    using Gauss-Hermite quadrature.

    :param f: integrand function. Takes one input of shape ?xD.
    :param H: Number of Gauss-Hermite evaluation points.
    :param Din: Number of input dimensions. Needs to be known at call-time.
    :param Dout: Number of output dimensions. Defaults to (). Dout is assumed
        to leave out the item index, i.e. f actually maps (?xD)->(?x*Dout).
    :return: quadratures
    """
    # Figure out input shape information
    if Din is None:
        Din = means.shape[1]

    if Din is None:
        raise ValueError(
            "If `Din` is passed as `None`, `means` must have a known shape. "
            "Running mvnquad in `autoflow` without specifying `Din` and `Dout` "
            "is problematic. Consider using your own session."
        )  # pragma: no cover

    xn, wn = mvhermgauss(H, Din)
    N = means.shape[0]

    # transform points based on Gaussian parameters
    cholXcov = tf.linalg.cholesky(covs)  # NxDxD
    Xt = tf.linalg.matmul(
        cholXcov, tf.tile(xn[None, :, :], (N, 1, 1)), transpose_b=True
    )  # NxDxH**D
    X = 2.0 ** 0.5 * Xt + tf.expand_dims(means, 2)  # NxDxH**D
    Xr = tf.reshape(tf.transpose(X, [2, 0, 1]), (-1, Din))  # (H**D*N)xD

    # perform quadrature
    fevals = func(Xr)
    if Dout is None:
        Dout = tuple((d if type(d) is int else d.value) for d in fevals.shape[1:])

    if any([d is None for d in Dout]):
        raise ValueError(
            "If `Dout` is passed as `None`, the output of `func` must have known "
            "shape. Running mvnquad in `autoflow` without specifying `Din` and `Dout` "
            "is problematic. Consider using your own session."
        )  # pragma: no cover
    fX = tf.reshape(
        fevals,
        (
            H ** Din,
            N,
        )
        + Dout,
    )
    wr = np.reshape(wn * np.pi ** (-Din * 0.5), (-1,) + (1,) * (1 + len(Dout)))
    return tf.reduce_sum(fX * wr, 0)


@check_shapes(
    "Fmu: [broadcast Din, N...]",
    "Fvar: [broadcast Din, N...]",
    "Ys.values(): [N...]",
    "return: [broadcast Dout, N...]",
)
def ndiagquad(
    funcs: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
    H: int,
    Fmu: Union[TensorType, Tuple[TensorType, ...], List[TensorType]],
    Fvar: Union[TensorType, Tuple[TensorType, ...], List[TensorType]],
    logspace: bool = False,
    **Ys: TensorType,
) -> tf.Tensor:
    """
    Computes N Gaussian expectation integrals of one or more functions
    using Gauss-Hermite quadrature. The Gaussians must be independent.

    The means and variances of the Gaussians are specified by Fmu and Fvar.
    The N-integrals are assumed to be taken wrt the last dimensions of Fmu, Fvar.

    `Fmu`, `Fvar`, `Ys` should all have same shape, with overall size `N`.

    :param funcs: the integrand(s):
        Callable or Iterable of Callables that operates elementwise
    :param H: number of Gauss-Hermite quadrature points
    :param Fmu: array/tensor or `Din`-tuple/list thereof
    :param Fvar: array/tensor or `Din`-tuple/list thereof
    :param logspace: if True, funcs are the log-integrands and this calculates
        the log-expectation of exp(funcs)
    :param Ys: arrays/tensors; deterministic arguments to be passed by name
    :return: shape is the same as that of the first Fmu
    """
    warnings.warn(
        "Please use gpflow.quadrature.NDiagGHQuadrature instead "
        "(note the changed convention of how multi-dimensional quadrature is handled)",
        DeprecationWarning,
    )

    n_gh = H
    if isinstance(Fmu, (tuple, list)):
        dim = len(Fmu)
        shape = tf.shape(Fmu[0])
        Fmu = tf.stack(Fmu, axis=-1)
        Fvar = tf.stack(Fvar, axis=-1)
    else:
        dim = 1
        shape = tf.shape(Fmu)

    Fmu = tf.reshape(Fmu, (-1, dim))
    Fvar = tf.reshape(Fvar, (-1, dim))

    Ys = {Yname: tf.reshape(Y, (-1, 1)) for Yname, Y in Ys.items()}

    def wrapper(old_fun: Callable[..., tf.Tensor]) -> Callable[..., tf.Tensor]:
        @wraps(old_fun)
        def new_fun(X: TensorType, **Ys: TensorType) -> tf.Tensor:
            Xs = tf.unstack(tf.expand_dims(X, axis=-2), axis=-1)
            fun_eval = old_fun(*Xs, **Ys)
            return tf.cond(
                pred=tf.less(tf.rank(fun_eval), tf.rank(X)),
                true_fn=lambda: fun_eval[..., tf.newaxis],
                false_fn=lambda: fun_eval,
            )

        return new_fun

    if isinstance(funcs, Iterable):
        funcs = [wrapper(f) for f in funcs]
    else:
        funcs = wrapper(funcs)

    quadrature = NDiagGHQuadrature(dim, n_gh)
    if logspace:
        result = quadrature.logspace(funcs, Fmu, Fvar, **Ys)
    else:
        result = quadrature(funcs, Fmu, Fvar, **Ys)

    if isinstance(result, list):
        result = [tf.reshape(r, shape) for r in result]
    else:
        result = tf.reshape(result, shape)

    return result


@check_shapes(
    "Fmu: [N, Din]",
    "Fvar: [N, Din]",
    "Ys.values(): [broadcast N, .]",
    "return: [broadcast n_funs, N, P]",
)
def ndiag_mc(
    funcs: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
    S: int,
    Fmu: TensorType,
    Fvar: TensorType,
    logspace: bool = False,
    epsilon: Optional[TensorType] = None,
    **Ys: TensorType,
) -> tf.Tensor:
    """
    Computes N Gaussian expectation integrals of one or more functions
    using Monte Carlo samples. The Gaussians must be independent.

    `Fmu`, `Fvar`, `Ys` should all have same shape, with overall size `N`.

    :param funcs: the integrand(s):
        Callable or Iterable of Callables that operates elementwise
    :param S: number of Monte Carlo sampling points
    :param Fmu: array/tensor
    :param Fvar: array/tensor
    :param logspace: if True, funcs are the log-integrands and this calculates
        the log-expectation of exp(funcs)
    :param Ys: arrays/tensors; deterministic arguments to be passed by name
    :return: shape is the same as that of the first Fmu
    """
    N, D = tf.shape(Fmu)[0], tf.shape(Fvar)[1]

    if epsilon is None:
        epsilon = tf.random.normal(shape=[S, N, D], dtype=default_float())

    mc_x = Fmu[None, :, :] + tf.sqrt(Fvar[None, :, :]) * epsilon
    mc_Xr = tf.reshape(mc_x, (S * N, D))

    for name, Y in Ys.items():
        D_out = Y.shape[1]
        # we can't rely on broadcasting and need tiling
        mc_Yr = tf.tile(Y[None, ...], [S, 1, 1])  # [S, N, D]_out
        Ys[name] = tf.reshape(mc_Yr, (S * N, D_out))  # S * [N, _]out

    def eval_func(func: Callable[..., tf.Tensor]) -> tf.Tensor:
        feval = func(mc_Xr, **Ys)
        feval = tf.reshape(feval, (S, N, -1))
        if logspace:
            log_S = tf.math.log(to_default_float(S))
            return tf.reduce_logsumexp(feval, axis=0) - log_S  # [N, D]
        else:
            return tf.reduce_mean(feval, axis=0)

    if isinstance(funcs, Iterable):
        return [eval_func(f) for f in funcs]
    else:
        return eval_func(funcs)


# gpflow/quadrature/gauss_hermite.py

# Copyright 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 typing import Sequence, Tuple

import numpy as np
import tensorflow as tf

from ..base import TensorType
from ..config import default_float
from ..experimental.check_shapes import check_shapes, inherit_check_shapes
from .base import GaussianQuadrature


@check_shapes(
    "return[0]: [N]",
    "return[1]: [N]",
)
def gh_points_and_weights(n_gh: int) -> Tuple[tf.Tensor, tf.Tensor]:
    r"""
    Given the number of Gauss-Hermite points n_gh,
    returns the points z and the weights dz to perform the following
    uni-dimensional gaussian quadrature:

    X ~ N(mean, stddev²)
    E[f(X)] = ∫ f(x) p(x) dx = \sum_{i=1}^{n_gh} f(mean + stddev*z_i) dz_i

    :param n_gh: Number of Gauss-Hermite points
    :returns: Points z and weights dz to compute uni-dimensional gaussian expectation
    """
    z, dz = np.polynomial.hermite.hermgauss(n_gh)
    z = z * np.sqrt(2)
    dz = dz / np.sqrt(np.pi)
    z, dz = z.astype(default_float()), dz.astype(default_float())
    return tf.convert_to_tensor(z), tf.convert_to_tensor(dz)


@check_shapes(
    "xs[all]: [.]",
    "return: [N_product, D]",
)
def list_to_flat_grid(xs: Sequence[TensorType]) -> tf.Tensor:
    """
    :param xs: List with d rank-1 Tensors, with shapes N1, N2, ..., Nd
    :return: Tensor with shape [N1*N2*...*Nd, d] representing the flattened
        d-dimensional grid built from the input tensors xs
    """
    return tf.reshape(tf.stack(tf.meshgrid(*xs), axis=-1), (-1, len(xs)))


@check_shapes(
    "zs[all]: [.]",
    "dzs[all]: [.]",
    "return[0]: [N_product, D]",
    "return[1]: [N_product, 1]",
)
def reshape_Z_dZ(
    zs: Sequence[TensorType], dzs: Sequence[TensorType]
) -> Tuple[tf.Tensor, tf.Tensor]:
    """
    :param zs: List with d rank-1 Tensors, with shapes N1, N2, ..., Nd
    :param dzs: List with d rank-1 Tensors, with shapes N1, N2, ..., Nd
    :returns: points Z, Tensor with shape [N1*N2*...*Nd, D],
        and weights dZ, Tensor with shape [N1*N2*...*Nd, 1]
    """
    Z = list_to_flat_grid(zs)
    dZ = tf.reduce_prod(list_to_flat_grid(dzs), axis=-1, keepdims=True)
    return Z, dZ


@check_shapes(
    "x: [batch...]",
    "return: [n, batch...]",
)
def repeat_as_list(x: TensorType, n: int) -> Sequence[tf.Tensor]:
    """
    :param x: Array/Tensor to be repeated
    :param n: Integer with the number of repetitions
    :return: List of n repetitions of Tensor x
    """
    return [x for _ in range(n)]


@check_shapes(
    "return[0]: [n_quad_points, D]",
    "return[1]: [n_quad_points, 1]",
)
def ndgh_points_and_weights(dim: int, n_gh: int) -> Tuple[tf.Tensor, tf.Tensor]:
    r"""
    :param dim: dimension of the multivariate normal
    :param n_gh: number of Gauss-Hermite points per dimension
    :returns: points Z, Tensor with shape [n_gh**dim, D],
        and weights dZ, Tensor with shape [n_gh**dim, 1]
    """
    z, dz = gh_points_and_weights(n_gh)
    zs = repeat_as_list(z, dim)
    dzs = repeat_as_list(dz, dim)
    return reshape_Z_dZ(zs, dzs)


class NDiagGHQuadrature(GaussianQuadrature):
    def __init__(self, dim: int, n_gh: int) -> None:
        """
        :param dim: dimension of the multivariate normal
        :param n_gh: number of Gauss-Hermite points per dimension
        """
        self.dim = dim
        self.n_gh = n_gh
        self.n_gh_total = n_gh ** dim
        Z, dZ = ndgh_points_and_weights(self.dim, self.n_gh)
        self.Z = tf.ensure_shape(Z, (self.n_gh_total, self.dim))
        self.dZ = tf.ensure_shape(dZ, (self.n_gh_total, 1))

    @inherit_check_shapes
    def _build_X_W(self, mean: TensorType, var: TensorType) -> Tuple[tf.Tensor, tf.Tensor]:
        """
        :param mean: Array/Tensor with shape [b1, b2, ..., bX, dim], usually [N, dim],
            representing the mean of a dim-Variate Gaussian distribution
        :param var: Array/Tensor with shape b1, b2, ..., bX, dim], usually [N, dim],
            representing the variance of a dim-Variate Gaussian distribution
        :return: points X, Tensor with shape [n_gh_total, b1, b2, ..., bX, dim],
            usually [n_gh_total, N, dim],
            and weights W, a Tensor with shape [n_gh_total, b1, b2, ..., bX, 1],
            usually [n_gh_total, N, 1]
        """

        batch_shape_broadcast = tf.ones(tf.rank(mean) - 1, dtype=tf.int32)
        shape_aux = tf.concat([[self.n_gh_total], batch_shape_broadcast], axis=0)

        # mean, var: [b1, b2, ..., bX, dim], usually [N, dim]
        mean = tf.expand_dims(mean, 0)
        stddev = tf.expand_dims(tf.sqrt(var), 0)
        # mean, stddev: [1, b1, b2, ..., bX, dim], usually [1, N, dim]

        Z = tf.cast(tf.reshape(self.Z, tf.concat([shape_aux, [self.dim]], axis=0)), mean.dtype)
        dZ = tf.cast(tf.reshape(self.dZ, tf.concat([shape_aux, [1]], axis=0)), mean.dtype)

        X = mean + stddev * Z
        W = dZ
        # X: [n_gh_total, b1, b2, ..., bX, dim], usually [n_gh_total, N, dim]
        # W: [n_gh_total,  1,  1, ...,  1,   1], usually [n_gh_total, N,   1]

        return X, W


# gpflow/utilities/__init__.py

from .bijectors import positive, triangular, triangular_size
from .misc import is_variable, set_trainable, to_default_float, to_default_int, training_loop
from .model_utils import add_likelihood_noise_cov, add_noise_cov, assert_params_false
from .multipledispatch import Dispatcher
from .traversal import (
    deepcopy,
    freeze,
    leaf_components,
    multiple_assign,
    parameter_dict,
    print_summary,
    read_values,
    reset_cache_bijectors,
    select_dict_parameters_with_prior,
    tabulate_module_summary,
    traverse_module,
)

__all__ = [
    "Dispatcher",
    "add_likelihood_noise_cov",
    "add_noise_cov",
    "assert_params_false",
    "bijectors",
    "deepcopy",
    "freeze",
    "is_variable",
    "leaf_components",
    "misc",
    "model_utils",
    "multiple_assign",
    "multipledispatch",
    "ops",
    "parameter_dict",
    "parameter_or_function",
    "positive",
    "print_summary",
    "read_values",
    "reset_cache_bijectors",
    "select_dict_parameters_with_prior",
    "set_trainable",
    "tabulate_module_summary",
    "to_default_float",
    "to_default_int",
    "training_loop",
    "traversal",
    "traverse_module",
    "triangular",
    "triangular_size",
]


# gpflow/utilities/bijectors.py

# Copyright 2019-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 typing import Optional

import tensorflow as tf
import tensorflow_probability as tfp

from .. import config
from ..experimental.check_shapes import check_shapes
from .misc import to_default_float

__all__ = ["positive", "triangular", "triangular_size"]


def positive(lower: Optional[float] = None, base: Optional[str] = None) -> tfp.bijectors.Bijector:
    """
    Returns a positive bijector (a reversible transformation from real to positive numbers).

    :param lower: overrides default lower bound
        (if None, defaults to gpflow.config.default_positive_minimum())
    :param base: overrides base positive bijector
        (if None, defaults to gpflow.config.default_positive_bijector())
    :returns: a bijector instance
    """
    bijector = base if base is not None else config.default_positive_bijector()
    bijector = config.positive_bijector_type_map()[bijector.lower()]()

    lower_bound = lower if lower is not None else config.default_positive_minimum()

    if lower_bound != 0.0:
        shift = tfp.bijectors.Shift(to_default_float(lower_bound))
        bijector = tfp.bijectors.Chain([shift, bijector])  # from unconstrained to constrained
    return bijector


def triangular() -> tfp.bijectors.Bijector:
    """
    Returns instance of a (lower) triangular bijector.
    """
    return tfp.bijectors.FillTriangular()


@check_shapes(
    "n: []",
    "return: []",
)
def triangular_size(n: tf.Tensor) -> tf.Tensor:
    """
    Returns the number of non-zero elements in an `n` by `n` triangular matrix.
    """
    return n * (n + 1) // 2


# gpflow/utilities/misc.py

# Copyright 2017-2021 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 typing import Callable, Iterable, List, Optional, Union

import tensorflow as tf
import tensorflow_probability as tfp

from ..base import TensorData
from ..config import default_float, default_int
from ..experimental.check_shapes import check_shapes
from .ops import cast

__all__ = [
    "is_variable",
    "set_trainable",
    "to_default_float",
    "to_default_int",
    "training_loop",
]


@check_shapes(
    "x: [any...]",
    "return: [any...]",
)
def to_default_int(x: TensorData) -> tf.Tensor:
    return cast(x, dtype=default_int())


@check_shapes(
    "x: [any...]",
    "return: [any...]",
)
def to_default_float(x: TensorData) -> tf.Tensor:
    return cast(x, dtype=default_float())


def set_trainable(model: Union[tf.Module, Iterable[tf.Module]], flag: bool) -> None:
    """
    Set trainable flag for all :class:`tf.Variable`\ s and :class:`gpflow.Parameter`\ s in a
    :class:`tf.Module` or collection of :class:`tf.Module`\ s.
    """
    modules = [model] if isinstance(model, tf.Module) else model

    for mod in modules:
        for variable in mod.variables:
            variable._trainable = flag


def is_variable(t: TensorData) -> bool:
    """
    Returns whether the `t` is a TensorFlow variable.
    """
    return isinstance(t, (tf.Variable, tfp.util.TransformedVariable))


def training_loop(
    closure: Callable[[], tf.Tensor],
    optimizer: Optional[tf.optimizers.Optimizer] = None,
    var_list: Optional[List[tf.Variable]] = None,
    maxiter: int = 1_000,
    compile: bool = False,
) -> None:
    """
    Simple generic training loop. At each iteration uses a GradientTape to compute
    the gradients of a loss function with respect to a set of variables.

    :param closure: Callable that constructs a loss function based on data and model being trained
    :param optimizer: tf.optimizers or tf.keras.optimizers that updates variables by applying the
        corresponding loss gradients. Adam is a default optimizer with default settings.
    :param var_list: List of model variables to be learnt during training
    :param maxiter: Maximum number of
    :return:
    """

    safe_optimizer = tf.optimizers.Adam() if optimizer is None else optimizer
    safe_var_list = [] if var_list is None else var_list

    def optimization_step() -> None:
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(safe_var_list)
            loss = closure()
        grads = tape.gradient(loss, safe_var_list)
        safe_optimizer.apply_gradients(zip(grads, safe_var_list))

    if compile:
        optimization_step = tf.function(optimization_step)

    for _ in range(maxiter):
        optimization_step()


# gpflow/utilities/model_utils.py

from typing import Any, Callable

import tensorflow as tf

from ..base import TensorType
from ..experimental.check_shapes import check_shapes
from ..likelihoods import Gaussian


def assert_params_false(
    called_method: Callable[..., Any],
    **kwargs: bool,
) -> None:
    """
    Asserts that parameters are ``False``.

    :param called_method: The method or function that is calling this. Used for nice error messages.
    :param kwargs: Parameters that must be ``False``.
    :raises NotImplementedError: If any ``kwargs`` are ``True``.
    """
    errors_str = ", ".join(f"{param}={value}" for param, value in kwargs.items() if value)
    if errors_str:
        raise NotImplementedError(
            f"{called_method.__qualname__} does not currently support: {errors_str}"
        )


@check_shapes(
    "K: [batch..., N, N]",
    "likelihood_variance: [broadcast batch..., broadcast N]",
    "return: [batch..., N, N]",
)
def add_noise_cov(K: tf.Tensor, likelihood_variance: TensorType) -> tf.Tensor:
    """
    Returns K + σ², where σ² is the diagonal likelihood noise variance.
    """
    k_diag = tf.linalg.diag_part(K)
    return tf.linalg.set_diag(K, k_diag + likelihood_variance)


@check_shapes(
    "K: [batch..., N, N]",
    "X: [batch..., N, D]",
    "return: [batch..., N, N]",
)
def add_likelihood_noise_cov(K: tf.Tensor, likelihood: Gaussian, X: TensorType) -> tf.Tensor:
    """
    Returns K + σ², where σ² is the likelihood noise variance.
    """
    return add_noise_cov(K, tf.squeeze(likelihood.variance_at(X), axis=-1))


# gpflow/utilities/multipledispatch.py

# Copyright 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 typing import Any, Callable, Optional, Tuple, Type, TypeVar, Union

from multipledispatch import Dispatcher as GeneratorDispatcher
from multipledispatch.dispatcher import str_signature, variadic_signature_matches
from multipledispatch.variadic import isvariadic

__all__ = ["Dispatcher"]


AnyCallable = Callable[..., Any]
_C = TypeVar("_C", bound=AnyCallable)
Types = Union[Type[Any], Tuple[Type[Any], ...]]


class Dispatcher(GeneratorDispatcher):
    """
    multipledispatch.Dispatcher uses a generator to yield the
    desired function implementation, which is problematic as TensorFlow's
    autograph is not able to compile code that passes through generators.

    This class overwrites the problematic method in the original
    Dispatcher and solely makes use of simple for-loops, which are
    compilable by AutoGraph.
    """

    def register(self, *types: Types, **kwargs: Any) -> Callable[[_C], _C]:
        # Override to add type hints...
        result: Callable[[_C], _C] = super().register(*types, **kwargs)
        return result

    def dispatch(self, *types: Types) -> Optional[AnyCallable]:
        """
        Returns matching function for `types`; if not existing returns None.
        """
        if types in self.funcs:
            result: AnyCallable = self.funcs[types]
            return result

        return self.get_first_occurrence(*types)

    def dispatch_or_raise(self, *types: Types) -> AnyCallable:
        """
        Returns matching function for `types`; if not existing raises an error.
        """
        f = self.dispatch(*types)
        if f is None:
            raise NotImplementedError(
                f"Could not find signature for {self.name}: <{str_signature(types)}>"
            )
        return f

    def get_first_occurrence(self, *types: Types) -> Optional[AnyCallable]:
        """
        Returns the first occurrence of a matching function

        Based on `multipledispatch.Dispatcher.dispatch_iter`, which
        returns an iterator of matching functions. This method uses
        the same logic to select functions, but simply returns the first
        element of the iterator. If no matching functions are found,
        `None` is returned.
        """
        n = len(types)
        for signature in self.ordering:
            if len(signature) == n and all(map(issubclass, types, signature)):  # type: ignore[arg-type]
                result: AnyCallable = self.funcs[signature]
                return result
            elif len(signature) and isvariadic(signature[-1]):
                if variadic_signature_matches(types, signature):
                    result = self.funcs[signature]
                    return result
        return None


# gpflow/utilities/ops.py

# 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.

import copy
from typing import Any, Callable, List, Optional, Union

import tensorflow as tf
import tensorflow_probability as tfp

from ..base import AnyNDArray
from ..experimental.check_shapes import check_shapes


@check_shapes(
    "value: [any...]",
    "return: [any...]",
)
def cast(
    value: Union[tf.Tensor, AnyNDArray], dtype: tf.DType, name: Optional[str] = None
) -> tf.Tensor:
    if not tf.is_tensor(value):
        # TODO(awav): Release TF2.2 resolves this issue
        # workaround for https://github.com/tensorflow/tensorflow/issues/35938
        return tf.convert_to_tensor(value, dtype, name=name)
    return tf.cast(value, dtype, name=name)


@check_shapes(
    "value: []",
    "return: [N, N]",
)
def eye(num: int, value: tf.Tensor, dtype: Optional[tf.DType] = None) -> tf.Tensor:
    if dtype is not None:
        value = cast(value, dtype)
    return tf.linalg.diag(tf.fill([num], value))


@check_shapes(
    "tensor: [any...]",
    "return: [transposed_any...]",
)
def leading_transpose(tensor: tf.Tensor, perm: List[Any], leading_dim: int = 0) -> tf.Tensor:
    """
    Transposes tensors with leading dimensions.

    Leading dimensions in permutation list represented via ellipsis `...` and is of type
    List[Union[int, type(...)]  (please note, due to mypy issues, List[Any] is used instead).  When
    leading dimensions are found, `transpose` method considers them as a single grouped element
    indexed by 0 in `perm` list. So, passing `perm=[-2, ..., -1]`, you assume that your input tensor
    has [..., A, B] shape, and you want to move leading dims between A and B dimensions.  Dimension
    indices in permutation list can be negative or positive. Valid positive indices start from 1 up
    to the tensor rank, viewing leading dimensions `...` as zero index.

    Example::

        a = tf.random.normal((1, 2, 3, 4, 5, 6))
        # [..., A, B, C],
        # where A is 1st element,
        # B is 2nd element and
        # C is 3rd element in
        # permutation list,
        # leading dimensions are [1, 2, 3]
        # which are 0th element in permutation list
        b = leading_transpose(a, [3, -3, ..., -2])  # [C, A, ..., B]
        sess.run(b).shape

        output> (6, 4, 1, 2, 3, 5)

    :param tensor: TensorFlow tensor.
    :param perm: List of permutation indices.
    :returns: TensorFlow tensor.
    :raises ValueError: when `...` cannot be found.

    """
    perm = copy.copy(perm)
    idx = perm.index(...)
    perm[idx] = leading_dim

    rank = tf.rank(tensor)
    perm_tf = perm % rank

    leading_dims = tf.range(rank - len(perm) + 1)
    perm = tf.concat([perm_tf[:idx], leading_dims, perm_tf[idx + 1 :]], 0)
    return tf.transpose(tensor, perm)


@check_shapes(
    "a: [a_shape...]",
    "b: [b_shape...]",
    "return: [a_shape..., b_shape...]",
)
def broadcasting_elementwise(
    op: Callable[[tf.Tensor, tf.Tensor], tf.Tensor], a: tf.Tensor, b: tf.Tensor
) -> tf.Tensor:
    """
    Apply binary operation `op` to every pair in tensors `a` and `b`.

    :param op: binary operator on tensors, e.g. tf.add, tf.substract
    """
    flatres = op(tf.reshape(a, [-1, 1]), tf.reshape(b, [1, -1]))
    return tf.reshape(flatres, tf.concat([tf.shape(a), tf.shape(b)], 0))


@check_shapes(
    "X: [batch..., N, D]",
    "X2: [batch2..., N2, D]",
    "return: [batch..., N, batch2..., N2] if X2 is not None",
    "return: [batch..., N, N] if X2 is None",
)
def square_distance(X: tf.Tensor, X2: Optional[tf.Tensor]) -> tf.Tensor:
    """
    Returns ||X - X2ᵀ||²

    Due to the implementation and floating-point imprecision, the
    result may actually be very slightly negative for entries very
    close to each other.
    """
    if X2 is None:
        Xs = tf.reduce_sum(tf.square(X), axis=-1, keepdims=True)
        dist = -2 * tf.matmul(X, X, transpose_b=True)
        dist += Xs + tf.linalg.adjoint(Xs)
        return dist
    Xs = tf.reduce_sum(tf.square(X), axis=-1)
    X2s = tf.reduce_sum(tf.square(X2), axis=-1)
    dist = -2 * tf.tensordot(X, X2, [[-1], [-1]])
    dist += broadcasting_elementwise(tf.add, Xs, X2s)
    return dist


@check_shapes(
    "X: [batch..., N, D]",
    "X2: [batch2..., N2, D]",
    "return: [batch..., N, batch2..., N2, D] if X2 is not None",
    "return: [batch..., N, N, D] if X2 is None",
)
def difference_matrix(X: tf.Tensor, X2: Optional[tf.Tensor]) -> tf.Tensor:
    """
    Returns (X - X2ᵀ)
    """
    if X2 is None:
        X2 = X
        diff = X[..., :, tf.newaxis, :] - X2[..., tf.newaxis, :, :]
        return diff
    Xshape = tf.shape(X)
    X2shape = tf.shape(X2)
    X = tf.reshape(X, (-1, Xshape[-1]))
    X2 = tf.reshape(X2, (-1, X2shape[-1]))
    diff = X[:, tf.newaxis, :] - X2[tf.newaxis, :, :]
    diff = tf.reshape(diff, tf.concat((Xshape[:-1], X2shape[:-1], [Xshape[-1]]), 0))
    return diff


@check_shapes(
    "X: [N, D]",
    "latent_dim: []",
    "return: [N, Q]",
)
def pca_reduce(X: tf.Tensor, latent_dim: tf.Tensor) -> tf.Tensor:
    """
    Linearly reduce the dimensionality of the input points `X` to `latent_dim` dimensions.

    :param X: Data to reduce.
    :param latent_dim: Number of latent dimension, Q < D.
    :return: PCA projection array.

    """
    if latent_dim > X.shape[1]:  # pragma: no cover
        raise ValueError("Cannot have more latent dimensions than observed")
    X_cov = tfp.stats.covariance(X)
    evals, evecs = tf.linalg.eigh(X_cov)
    W = evecs[:, -latent_dim:]
    return (X - tf.reduce_mean(X, axis=0, keepdims=True)) @ W


# gpflow/utilities/parameter_or_function.py

# Copyright 2022 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 typing import Optional, Union

import tensorflow as tf

from ..base import Parameter, TensorData, TensorType
from ..experimental.check_shapes import check_shapes
from ..functions import Function
from . import positive

ConstantOrFunction = Union[Function, TensorData]
ParameterOrFunction = Union[Function, Parameter]


def prepare_parameter_or_function(
    value: ConstantOrFunction,
    *,
    lower_bound: Optional[float] = None,
) -> ParameterOrFunction:
    if isinstance(value, Function):
        return value
    else:
        if lower_bound is None:
            return Parameter(value)
        else:
            return Parameter(value, transform=positive(lower_bound))


@check_shapes(
    "X: [batch..., N, D]",
    "return: [broadcast batch..., broadcast N, broadcast P]",
)
def evaluate_parameter_or_function(
    value: ParameterOrFunction,
    X: TensorType,
    *,
    lower_bound: Optional[float] = None,
) -> TensorType:
    if isinstance(value, Function):
        result = value(X)
        if lower_bound is not None:
            result = tf.maximum(result, lower_bound)
        return result
    else:
        return value


# gpflow/utilities/traversal.py

# Copyright 2017-2021 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.

import copy
import re
from functools import lru_cache
from typing import Any, Callable, Dict, Mapping, Optional, Pattern, Tuple, Type, TypeVar, Union

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from packaging.version import Version
from tabulate import tabulate
from tensorflow.python.util.object_identity import Reference

from ..base import AnyNDArray, Parameter
from ..config import default_summary_fmt

__all__ = [
    "deepcopy",
    "freeze",
    "leaf_components",
    "multiple_assign",
    "parameter_dict",
    "print_summary",
    "read_values",
    "reset_cache_bijectors",
    "select_dict_parameters_with_prior",
    "tabulate_module_summary",
]
LeafComponent = Union[tf.Variable, tf.Tensor, Parameter]
LeafVariable = Union[tf.Variable, Parameter]
HashableTensor = Union[Reference, Parameter]
TraverseInput = TypeVar("TraverseInput", tf.Variable, tf.Module, Parameter)
State = TypeVar("State")
Path = str
Accumulator = Tuple[Path, State]
TraverseUpdateCallable = Callable[[TraverseInput, Path, State], State]


def multiple_assign(module: tf.Module, parameters: Mapping[Path, tf.Tensor]) -> None:
    """
    Multiple assign takes a dictionary with new values. Dictionary keys are paths to the
    `tf.Variable`s or `gpflow.Parameter` of the input module.

    :param module: `tf.Module`.
    :param parameters: a dictionary with keys of the form ".module.path.to.variable" and new value tensors.
    """
    reference_var_dict = parameter_dict(module)
    for path, value in parameters.items():
        reference_var_dict[path].assign(value)


def read_values(module: tf.Module) -> Dict[Path, AnyNDArray]:
    """Returns a dictionary of numpy values of the module parameters (variables)."""
    return {k: v.numpy() for k, v in parameter_dict(module).items()}


def parameter_dict(module: tf.Module) -> Dict[Path, LeafVariable]:
    """
    Returns a dictionary of parameters (variables) for the `tf.Module` component.
    Dictionary keys are relative paths to the attributes to which parameters (variables) assigned to.

        class SubModule(tf.Module):
            def __init__(self):
                self.parameter = gpflow.Parameter(1.0)
                self.variable = tf.Variable(1.0)

        class Module(tf.Module):
            def __init__(self):
                self.submodule = SubModule()

        m = Module()
        params = parameter_dict(m)
        # {
        #   ".submodule.parameter": <parameter object>,
        #   ".submodule.variable": <variable object>
        # }
    """
    param_dict = leaf_components(module)
    return {f".{key.split('.', 1)[-1]}": value for key, value in param_dict.items()}


def print_summary(module: tf.Module, fmt: Optional[str] = None) -> None:
    """
    Prints a summary of the parameters and variables contained in a tf.Module.
    """
    fmt = fmt if fmt is not None else default_summary_fmt()
    if fmt == "notebook":
        from IPython.core.display import HTML, display

        tab = tabulate_module_summary(module, "html")
        display(HTML(tab))
    else:
        print(tabulate_module_summary(module, fmt))


def tabulate_module_summary(module: tf.Module, tablefmt: Optional[str] = None) -> str:
    def get_transform(path: Path, var: LeafComponent) -> Optional[str]:
        if hasattr(var, "transform") and var.transform is not None:
            if isinstance(var.transform, tfp.bijectors.Chain):
                return " + ".join(b.__class__.__name__ for b in var.transform.bijectors[::-1])
            return var.transform.__class__.__name__  # type: ignore[no-any-return]
        return None

    def get_prior(path: Path, var: LeafComponent) -> Optional[str]:
        if hasattr(var, "prior") and var.prior is not None:
            return var.prior.name  # type: ignore[no-any-return]
        return None

    # list of (column_name: str, column_getter: Callable[[tf.Variable], str]) tuples:
    column_definition = [
        ("name", lambda path, var: path),
        ("class", lambda path, var: var.__class__.__name__),
        ("transform", get_transform),
        ("prior", get_prior),
        ("trainable", lambda path, var: var.trainable),
        ("shape", lambda path, var: var.shape),
        ("dtype", lambda path, var: var.dtype.name),
        ("value", lambda path, var: _str_tensor_value(var.numpy())),
    ]
    column_names, column_getters = zip(*column_definition)

    merged_leaf_components = _merge_leaf_components(leaf_components(module))

    column_values = [
        [getter(path, variable) for getter in column_getters]
        for path, variable in merged_leaf_components.items()
    ]
    # mypy claims it's wrong to pass a `None` tablefmt below. I think `tabulate` has bad type hints.
    return tabulate(column_values, headers=column_names, tablefmt=tablefmt)  # type: ignore[arg-type]


def leaf_components(input: tf.Module) -> Mapping[Path, LeafVariable]:
    return _get_leaf_components(input)


def _merge_leaf_components(input: Mapping[Path, LeafComponent]) -> Mapping[Path, LeafComponent]:

    ref_fn: Callable[[LeafComponent], HashableTensor] = lambda x: (
        x if isinstance(x, Parameter) else x.ref()
    )
    deref_fn: Callable[[HashableTensor], LeafComponent] = lambda x: (
        x if isinstance(x, Parameter) else x.deref()
    )

    input_values = {ref_fn(value) for value in input.values()}
    if len(input_values) == len(input):
        return input

    tmp_dict: Dict[HashableTensor, Path] = {}
    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: deref_fn(ref) for ref, key in tmp_dict.items()}


def _get_leaf_components(input_module: tf.Module) -> Mapping[Path, LeafVariable]:
    """
    Returns a list of tuples each corresponding to a gpflow.Parameter or tf.Variable in the each
    submodules of a given tf.Module. Each tuple consists of an specific Parameter (or Variable) and
    its relative path inside the module, which is constructed recursively by adding a prefix with
    the path to the current module. Designed to be used as a helper for the method 'print_summary'.

    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
    :return:
    """
    target_types = (Parameter, tf.Variable)
    input_name = input_module.__class__.__name__
    state: Dict[Path, LeafVariable] = {}
    accumulator = (input_name, state)

    def update_state(
        parameter_or_variable: LeafVariable, path: Path, state: Dict[Path, LeafVariable]
    ) -> Dict[Path, LeafVariable]:
        state[path] = parameter_or_variable
        return state

    state = traverse_module(input_module, accumulator, update_state, target_types)
    return state


def reset_cache_bijectors(input_module: tf.Module) -> tf.Module:
    """
    Recursively finds tfp.bijectors.Bijector-s inside the components of the tf.Module using `traverse_component`.
    Resets the caches stored inside each tfp.bijectors.Bijector.

    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
    :returns: same object but with all bijector caches reset
    """
    if Version(tfp.__version__) >= Version("0.11.0"):
        if hasattr(tfp.bijectors.Identity()._cache, "clear"):
            # implementation in `master` branch (checked 29 Sep 2020) provides clear():

            def _clear_bijector_cache(bijector: tfp.bijectors.Bijector) -> None:
                bijector._cache.clear()

        else:  # pragma: no cover
            # previous versions (including the versions 0.11.0 and 0.11.1 released as of 29 Sep 2020) provide reset(), but its implementation is broken

            def _clear_bijector_cache(bijector: tfp.bijectors.Bijector) -> None:
                # workaround for broken implementation of bijector._cache.reset():
                cache = bijector._cache
                cache_type = type(cache.forward)
                assert type(cache.inverse) == cache_type
                cache.__init__(cache.forward._func, cache.inverse._func, cache_type)

    else:  # pragma: no cover
        # fallback for backwards-compatibility with tensorflow_probability < 0.11.0

        def _clear_bijector_cache(bijector: tfp.bijectors.Bijector) -> None:
            # `_from_x` and `_from_y` are cache dictionaries for forward and inverse transformations
            bijector._from_x.clear()
            bijector._from_y.clear()

    target_types = (tfp.bijectors.Bijector,)
    accumulator = ("", None)

    def clear_bijector(bijector: tfp.bijectors.Bijector, _: Path, state: None) -> None:
        if not isinstance(bijector, tfp.bijectors.Bijector):
            return  # skip submodules that are not bijectors

        _clear_bijector_cache(bijector)

        if isinstance(bijector, tfp.bijectors.Chain):
            # recursively clear caches of sub-bijectors
            for m in bijector.submodules:
                if isinstance(m, tfp.bijectors.Bijector):
                    _clear_bijector_cache(m)

        return state

    traverse_module(input_module, accumulator, clear_bijector, target_types)
    return input_module


M = TypeVar("M", bound=tf.Module)


def deepcopy(input_module: M, memo: Optional[Dict[int, Any]] = None) -> M:
    """
    Returns a deepcopy of the input tf.Module. To do that first resets the caches stored inside each
    tfp.bijectors.Bijector to allow the deepcopy of the tf.Module.

    :param input_module: tf.Module including keras.Model, keras.layers.Layer and gpflow.Module.
    :param memo: passed through to func:`copy.deepcopy`
        (see https://docs.python.org/3/library/copy.html).
    :return: Returns a deepcopy of an input object.
    """
    return copy.deepcopy(reset_cache_bijectors(input_module), memo)  # type: ignore[no-any-return]


def freeze(input_module: M) -> M:
    """
    Returns a deepcopy of the input tf.Module with constants instead of variables and parameters.

    :param input_module: tf.Module or gpflow.Module.
    :return: Returns a frozen deepcopy of an input object.
    """
    objects_to_freeze = _get_leaf_components(input_module)
    memo_tensors = {id(v): tf.convert_to_tensor(v) for v in objects_to_freeze.values()}
    module_copy = deepcopy(input_module, memo_tensors)
    return module_copy


def traverse_module(
    m: TraverseInput,
    acc: Accumulator[State],
    update_cb: TraverseUpdateCallable[TraverseInput, State],
    target_types: Tuple[Type[Any], ...],
) -> State:
    """
    Recursively traverses `m`, accumulating in `acc` a path and a state until it finds an object of
    type in `target_types` to apply `update_cb` to update the accumulator `acc` and/or the object.

    :param m: tf.Module, tf.Variable or gpflow.Parameter
    :param acc: Tuple of path and state
    :param update_cb: Callable
    :param target_types: target class types
    :return:
    """
    path, state = acc

    new_state = state

    if isinstance(m, target_types):
        return update_cb(m, path, state)

    if isinstance(m, (list, tuple)):
        for term_idx, subterm in enumerate(m):
            new_acc = (f"{path}[{term_idx}]", new_state)
            new_state = traverse_module(subterm, new_acc, update_cb, target_types)
    elif isinstance(m, dict):
        for term_idx, subterm in m.items():
            new_acc = (f"{path}['{term_idx}']", new_state)
            new_state = traverse_module(subterm, new_acc, update_cb, target_types)
    elif isinstance(m, tf.Module):
        for name, submodule in vars(m).items():
            ignored_attributes = m._TF_MODULE_IGNORED_PROPERTIES
            # NOTE(awav): since tfp version 0.10.0, tfp.bijectors.Bijector instances have
            # `_parameters` dictionary with "self" references that cause
            # infinite recursive loop.
            if isinstance(m, tfp.bijectors.Bijector):
                ignored_attributes = ignored_attributes.union({"_parameters"})
            if name in ignored_attributes:
                continue
            new_acc = (f"{path}.{name}", new_state)
            new_state = traverse_module(submodule, new_acc, update_cb, target_types)
    return new_state


@lru_cache()
def _first_three_elements_regexp() -> Pattern[str]:
    num_re = r"[+\-]?(?:0|[1-9]\d*)(?:\.\d*)?(?:[eE][+\-]?\d+)?"
    pat_re = rf"^(?:(\[+)\s*)?({num_re})(?:\s+({num_re})(?:\s+({num_re}))?)?.*?"
    return re.compile(pat_re)


def _str_tensor_value(value: AnyNDArray) -> str:
    value_str = str(np.around(value, 5))
    if value.size <= 3:
        return value_str

    max_chars = 500
    value_str = value_str[:max_chars]
    regexp = _first_three_elements_regexp()
    match = regexp.match(value_str)
    assert match is not None
    brackets, elem1, elem2, elem3 = match.groups()

    out = f"{elem1}"
    if elem2 is not None:
        out = f"{out}{f', {elem2}'}"
        if elem3 is not None:
            out = f"{out}{f', {elem3}'}"
    if brackets is not None:
        out = f"{brackets}{out}..."

    return out


def select_dict_parameters_with_prior(model: tf.Module) -> Dict[Path, Parameter]:
    """Collects parameters with prior into a dictionary."""
    return {
        k: p
        for k, p in parameter_dict(model).items()
        if hasattr(p, "prior") and p.prior is not None
    }


# tests/gpflow/conditionals/test_broadcasted_conditionals.py

# Copyright 2017 the GPflow authors.
#
# 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.
"""
This test suite will check if the conditionals broadcast correctly
when the input tensors have leading dimensions.
"""

from typing import cast

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
import gpflow.inducing_variables.multioutput as mf
import gpflow.kernels.multioutput as mk
from gpflow.base import AnyNDArray, SamplesMeanAndVariance
from gpflow.conditionals import sample_conditional
from gpflow.conditionals.util import mix_latent_gp

# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


class Data:
    S1, S2, N, M = (
        7,  # num samples 1
        6,  # num samples 2
        4,  # num datapoints
        3,  # num inducing
    )
    Dx, Dy, L = (
        2,  # input dim
        5,  # output dim
        4,  # num latent GPs
    )
    W = np.random.randn(Dy, L)  # mixing matrix

    SX = np.random.randn(S1 * S2, N, Dx)
    S1_S2_X = np.reshape(SX, [S1, S2, N, Dx])

    Z = np.random.randn(M, Dx)


@pytest.mark.parametrize("full_cov", [False, True])
@pytest.mark.parametrize("white", [True, False])
@pytest.mark.parametrize("conditional_type", ["mixing", "Z", "inducing_points"])
def test_conditional_broadcasting(full_cov: bool, white: bool, conditional_type: str) -> None:
    """
    Test that the `conditional` and `sample_conditional` broadcasts correctly
    over leading dimensions of Xnew. Xnew can be shape [..., N, D],
    and conditional should broadcast over the [...].
    """
    q_mu = np.random.randn(Data.M, Data.Dy)
    q_sqrt: AnyNDArray = np.tril(np.random.randn(Data.Dy, Data.M, Data.M), -1)

    if conditional_type == "Z":
        inducing_variable = Data.Z
        kernel = gpflow.kernels.Matern52(lengthscales=0.5)
    elif conditional_type == "inducing_points":
        inducing_variable = gpflow.inducing_variables.InducingPoints(Data.Z)
        kernel = gpflow.kernels.Matern52(lengthscales=0.5)
    elif conditional_type == "mixing":
        # variational params have different output dim in this case
        q_mu = np.random.randn(Data.M, Data.L)
        q_sqrt = np.tril(np.random.randn(Data.L, Data.M, Data.M), -1)
        inducing_variable = mf.SharedIndependentInducingVariables(
            gpflow.inducing_variables.InducingPoints(Data.Z)
        )
        kernel = mk.LinearCoregionalization(
            kernels=[gpflow.kernels.Matern52(lengthscales=0.5) for _ in range(Data.L)],
            W=Data.W,
        )
    else:
        raise NotImplementedError

    if conditional_type == "mixing" and full_cov:
        pytest.skip("combination is not implemented")

    num_samples = 5

    def sample_conditional_fn(X: tf.Tensor) -> SamplesMeanAndVariance:
        return cast(
            SamplesMeanAndVariance,
            sample_conditional(
                X,
                inducing_variable,
                kernel,
                tf.convert_to_tensor(q_mu),
                q_sqrt=tf.convert_to_tensor(q_sqrt),
                white=white,
                full_cov=full_cov,
                num_samples=num_samples,
            ),
        )

    samples: AnyNDArray = np.array([sample_conditional_fn(X)[0] for X in Data.SX])
    means: AnyNDArray = np.array([sample_conditional_fn(X)[1] for X in Data.SX])
    variables: AnyNDArray = np.array([sample_conditional_fn(X)[2] for X in Data.SX])

    samples_S12, means_S12, vars_S12 = sample_conditional(
        Data.SX,
        inducing_variable,
        kernel,
        tf.convert_to_tensor(q_mu),
        q_sqrt=tf.convert_to_tensor(q_sqrt),
        white=white,
        full_cov=full_cov,
        num_samples=num_samples,
    )

    samples_S1_S2, means_S1_S2, vars_S1_S2 = sample_conditional(
        Data.S1_S2_X,
        inducing_variable,
        kernel,
        tf.convert_to_tensor(q_mu),
        q_sqrt=tf.convert_to_tensor(q_sqrt),
        white=white,
        full_cov=full_cov,
        num_samples=num_samples,
    )

    assert_allclose(samples_S12.shape, samples.shape)
    assert_allclose(samples_S1_S2.shape, [Data.S1, Data.S2, num_samples, Data.N, Data.Dy])
    assert_allclose(means_S12, means)
    assert_allclose(vars_S12, variables)
    assert_allclose(means_S1_S2.numpy().reshape(Data.S1 * Data.S2, Data.N, Data.Dy), means)
    if full_cov:
        vars_s1_s2 = vars_S1_S2.numpy().reshape(Data.S1 * Data.S2, Data.Dy, Data.N, Data.N)
        assert_allclose(vars_s1_s2, variables)
    else:
        vars_s1_s2 = vars_S1_S2.numpy().reshape(Data.S1 * Data.S2, Data.N, Data.Dy)
        assert_allclose(vars_s1_s2, variables)


# -------------------------------------------
# Test utility functions used in conditionals
# -------------------------------------------

# _mix_latent_gps
@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [True, False])
def test_broadcasting_mix_latent_gps(full_cov: bool, full_output_cov: bool) -> None:
    S, N = 7, 20  # batch size, num data points
    P, L = 10, 5  # observation dimensionality, num latent GPs
    W = np.random.randn(P, L)  # mixing matrix
    g_mu = np.random.randn(S, N, L)  # mean of the L latent GPs

    g_sqrt_diag: AnyNDArray = np.tril(np.random.randn(S * L, N, N), -1)  # [L*S, N, N]
    g_sqrt_diag = np.reshape(g_sqrt_diag, [L, S, N, N])
    g_var_diag: AnyNDArray = g_sqrt_diag @ np.transpose(g_sqrt_diag, [0, 1, 3, 2])  # [L, S, N, N]
    g_var = np.zeros([S, N, L, N, L])
    for l in range(L):
        g_var[:, :, l, :, l] = g_var_diag[l, :, :, :]  # replace diagonal elements by g_var_diag

    # reference numpy implementation for mean
    f_mu_ref: AnyNDArray = g_mu @ W.T  # [S, N, P]

    # reference numpy implementation for variance
    g_var_tmp = np.transpose(g_var, [0, 1, 3, 2, 4])  # [S, N, N, L, L]
    f_var_ref: AnyNDArray = W @ g_var_tmp @ W.T  # [S, N, N, P, P]
    f_var_ref = np.transpose(f_var_ref, [0, 1, 3, 2, 4])  # [S, N, P, N, P]

    if not full_cov:
        g_var_diag = np.array([g_var_diag[:, :, n, n] for n in range(N)])  # [N, L, S]
        g_var_diag = np.transpose(g_var_diag, [2, 0, 1])  # [S, N, L]

    # run gpflow's implementation
    f_mu, f_var = mix_latent_gp(
        tf.convert_to_tensor(W),
        tf.convert_to_tensor(g_mu),
        tf.convert_to_tensor(g_var_diag),
        full_cov,
        full_output_cov,
    )

    # we strip down f_var_ref to the elements we need
    if not full_output_cov and not full_cov:
        f_var_ref = np.array([f_var_ref[:, :, p, :, p] for p in range(P)])  # [P, S, N, N]
        f_var_ref = np.array([f_var_ref[:, :, n, n] for n in range(N)])  # [N, P, S]
        f_var_ref = np.transpose(f_var_ref, [2, 0, 1])  # [S, N, P]

    elif not full_output_cov and full_cov:
        f_var_ref = np.array([f_var_ref[:, :, p, :, p] for p in range(P)])  # [P, S, N, N]
        f_var_ref = np.transpose(f_var_ref, [1, 0, 2, 3])  # [S, P, N, N]

    elif full_output_cov and not full_cov:
        f_var_ref = np.array([f_var_ref[:, n, :, n, :] for n in range(N)])  # [N, S, P, P]
        f_var_ref = np.transpose(f_var_ref, [1, 0, 2, 3])  # [S, N, P, P]

    else:
        pass  # f_var_ref has shape [..., N, P, N, P] as expected

    # check equality for mean and variance of f
    assert_allclose(f_mu_ref, f_mu)
    assert_allclose(f_var_ref, f_var)


# tests/gpflow/conditionals/test_conditionals.py

# Copyright 2017 the GPflow authors.
#
# 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.

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
from gpflow import Parameter
from gpflow.base import AnyNDArray, MeanAndVariance
from gpflow.conditionals import conditional
from gpflow.config import default_float
from gpflow.kernels import Kernel
from gpflow.utilities.bijectors import triangular

rng = np.random.RandomState(123)

Ln = 2
Nn = 10
Mn = 20


@pytest.fixture(scope="module")
def kernel() -> Kernel:
    k = gpflow.kernels.Matern32() + gpflow.kernels.White()
    k.kernels[1].variance.assign(0.01)
    return k


@pytest.fixture(scope="module")
def Xdata() -> tf.Tensor:
    return tf.convert_to_tensor(rng.randn(Nn, 1))


@pytest.fixture(scope="module")
def Xnew() -> tf.Tensor:
    return tf.convert_to_tensor(rng.randn(Mn, 1))


@pytest.fixture(scope="module")
def mu() -> tf.Tensor:
    return tf.convert_to_tensor(rng.randn(Nn, Ln))


@pytest.fixture(scope="module")
def sqrt() -> tf.Tensor:
    return tf.convert_to_tensor(rng.randn(Nn, Ln))


@pytest.fixture(scope="module")
def chol(sqrt: tf.Tensor) -> tf.Tensor:
    return tf.stack([tf.linalg.diag(sqrt[:, i]) for i in range(Ln)])


@pytest.mark.parametrize("white", [True, False])
def test_diag(
    Xdata: tf.Tensor,
    Xnew: tf.Tensor,
    kernel: Kernel,
    mu: tf.Tensor,
    sqrt: tf.Tensor,
    chol: tf.Tensor,
    white: bool,
) -> None:
    Fstar_mean_1, Fstar_var_1 = conditional(Xnew, Xdata, kernel, mu, q_sqrt=sqrt, white=white)
    Fstar_mean_2, Fstar_var_2 = conditional(Xnew, Xdata, kernel, mu, q_sqrt=chol, white=white)

    mean_diff = Fstar_mean_1 - Fstar_mean_2
    var_diff = Fstar_var_1 - Fstar_var_2

    assert_allclose(mean_diff, 0)
    assert_allclose(var_diff, 0)


def test_whiten(
    Xdata: tf.Tensor, Xnew: tf.Tensor, kernel: Kernel, mu: tf.Tensor, sqrt: tf.Tensor
) -> None:
    """
    Make sure that predicting using the whitened representation is the
    sameas the non-whitened one.
    """

    K = kernel(Xdata) + tf.eye(Nn, dtype=default_float()) * 1e-6
    L = tf.linalg.cholesky(K)
    V = tf.linalg.triangular_solve(L, mu, lower=True)
    mean1, var1 = conditional(Xnew, Xdata, kernel, mu)
    mean2, var2 = conditional(Xnew, Xdata, kernel, V, white=True)

    assert_allclose(mean1, mean2)
    assert_allclose(var1, var2)


def test_gaussian_whiten(
    Xdata: tf.Tensor, Xnew: tf.Tensor, kernel: Kernel, mu: tf.Tensor, sqrt: tf.Tensor
) -> None:
    """
    Make sure that predicting using the whitened representation is the
    same as the non-whitened one.
    """
    F_sqrt = tf.convert_to_tensor(rng.rand(Nn, Ln))

    K = kernel(Xdata)
    L = tf.linalg.cholesky(K)
    V = tf.linalg.triangular_solve(L, mu, lower=True)
    V_prime = tf.linalg.diag(tf.transpose(F_sqrt))
    common_shape = tf.broadcast_static_shape(V_prime.shape, L.shape)
    L = tf.broadcast_to(L, common_shape)
    V_sqrt = tf.linalg.triangular_solve(L, tf.linalg.diag(tf.transpose(F_sqrt)), lower=True)

    Fstar_mean, Fstar_var = conditional(Xnew, Xdata, kernel, mu, q_sqrt=F_sqrt)
    Fstar_w_mean, Fstar_w_var = conditional(Xnew, Xdata, kernel, V, q_sqrt=V_sqrt, white=True)

    mean_diff = Fstar_w_mean - Fstar_mean
    var_diff = Fstar_w_var - Fstar_var

    assert_allclose(mean_diff, 0, atol=4)
    assert_allclose(var_diff, 0, atol=4)


@pytest.mark.parametrize("white", [True, False])
def test_q_sqrt_constraints(
    Xdata: tf.Tensor, Xnew: tf.Tensor, kernel: Kernel, mu: tf.Tensor, white: bool
) -> None:
    """Test that sending in an unconstrained q_sqrt returns the same conditional
    evaluation and gradients. This is important to match the behaviour of the KL, which
    enforces q_sqrt is triangular.
    """

    tril: AnyNDArray = np.tril(rng.randn(Ln, Nn, Nn))

    q_sqrt_constrained = Parameter(tril, transform=triangular())
    q_sqrt_unconstrained = Parameter(tril)

    diff_before_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy()
    assert_allclose(diff_before_gradient_step, 0)

    Fstars = []
    for q_sqrt in [q_sqrt_constrained, q_sqrt_unconstrained]:

        with tf.GradientTape() as tape:
            _, Fstar_var = conditional(Xnew, Xdata, kernel, mu, q_sqrt=q_sqrt, white=white)

        grad = tape.gradient(Fstar_var, q_sqrt.unconstrained_variable)
        q_sqrt.unconstrained_variable.assign_sub(grad)
        Fstars.append(Fstar_var)

    diff_Fstar_before_gradient_step = Fstars[0] - Fstars[1]
    assert_allclose(diff_Fstar_before_gradient_step, 0)

    diff_after_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy()
    assert_allclose(diff_after_gradient_step, 0)


@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("features_inducing_points", [False, True])
def test_base_conditional_vs_ref(full_cov: bool, features_inducing_points: bool) -> None:
    """
    Test that conditionals agree with a slow-but-clear numpy implementation
    """
    Dy, N, M, Dx = 5, 4, 3, 2
    X = np.random.randn(N, Dx)
    Z = np.random.randn(M, Dx)
    kern = gpflow.kernels.Matern52(lengthscales=0.5)
    q_mu = np.random.randn(M, Dy)
    q_sqrt: AnyNDArray = np.tril(np.random.randn(Dy, M, M), -1)

    def numpy_conditional(
        X: tf.Tensor, Z: tf.Tensor, kern: Kernel, q_mu: tf.Tensor, q_sqrt: tf.Tensor
    ) -> MeanAndVariance:
        Kmm = kern(Z, Z) + np.eye(M) * gpflow.config.default_jitter()
        Kmn = kern(Z, X)
        Knn = kern(X, X)

        Kmm, Kmn, Knn = [k.numpy() for k in [Kmm, Kmn, Knn]]
        Knm: AnyNDArray = Kmn.T

        Kmm, Kmn, Knm, Knn = [np.tile(k[None, :, :], [Dy, 1, 1]) for k in [Kmm, Kmn, Knm, Knn]]

        S = q_sqrt @ np.transpose(q_sqrt, [0, 2, 1])

        Kmm_inv = np.linalg.inv(Kmm)
        mean = np.einsum("dmn,dmM,Md->nd", Kmn, Kmm_inv, q_mu)
        cov = Knn + Knm @ Kmm_inv @ (S - Kmm) @ Kmm_inv @ Kmn
        return mean, cov

    mean_np, cov_np = numpy_conditional(X, Z, kern, q_mu, q_sqrt)

    if features_inducing_points:
        Z = gpflow.inducing_variables.InducingPoints(Z)

    mean_gpflow, cov_gpflow = [
        v.numpy()
        for v in gpflow.conditionals.conditional(
            X, Z, kern, q_mu, q_sqrt=tf.identity(q_sqrt), white=False, full_cov=full_cov
        )
    ]

    if not full_cov:
        cov_np = np.diagonal(cov_np, axis1=-1, axis2=-2).T

    assert_allclose(mean_np, mean_gpflow)
    assert_allclose(cov_np, cov_gpflow)


# tests/gpflow/conditionals/test_multioutput.py

from typing import Callable, List, Sequence, Tuple, cast

import numpy as np
import pytest
import scipy
import tensorflow as tf
from _pytest.fixtures import SubRequest

import gpflow
import gpflow.inducing_variables.multioutput as mf
import gpflow.kernels.multioutput as mk
from gpflow import set_trainable
from gpflow.base import AnyNDArray, RegressionData
from gpflow.conditionals import sample_conditional
from gpflow.conditionals.util import (
    fully_correlated_conditional,
    fully_correlated_conditional_repeat,
    independent_interdomain_conditional,
    sample_mvn,
)
from gpflow.config import default_float, default_jitter
from gpflow.experimental.check_shapes import ShapeChecker, check_shapes
from gpflow.inducing_variables import InducingPoints
from gpflow.kernels import SquaredExponential
from gpflow.likelihoods import Gaussian
from gpflow.models import SVGP

float_type = default_float()
rng = np.random.RandomState(99201)

# ------------------------------------------
# Helpers
# ------------------------------------------


@check_shapes(
    "Xnew: [N, D]",
    "return[0]: [n_models, N, P]",
    "return[1]: [n_models, N, P, N, P] if full_cov and full_output_cov",
    "return[1]: [n_models, N, P, P] if (not full_cov) and full_output_cov",
    "return[1]: [n_models, P, N, N] if full_cov and (not full_output_cov)",
    "return[1]: [n_models, N, P] if (not full_cov) and (not full_output_cov)",
)
def predict_all(
    models: Sequence[SVGP], Xnew: tf.Tensor, full_cov: bool, full_output_cov: bool
) -> Tuple[List[tf.Tensor], List[tf.Tensor]]:
    """
    Returns the mean and variance of f(Xnew) for each model in `models`.
    """
    ms, vs = [], []
    for model in models:
        m, v = model.predict_f(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        ms.append(m)
        vs.append(v)
    return ms, vs


@check_shapes(
    "arr[all]: [batch...]",
)
def assert_all_array_elements_almost_equal(arr: Sequence[tf.Tensor]) -> None:
    """
    Check if consecutive elements of `arr` are almost equal.
    """
    for i in range(len(arr) - 1):
        np.testing.assert_allclose(arr[i], arr[i + 1], atol=1e-5)


@check_shapes(
    "data[0]: [N, D]",
    "data[1]: [N, R]",
)
def check_equality_predictions(
    data: RegressionData, models: Sequence[SVGP], decimal: int = 3
) -> None:
    """
    Executes a couple of checks to compare the equality of predictions
    of different models. The models should be configured with the same
    training data (X, Y). The following checks are done:
    - check if elbo is (almost) equal for all models
    - check if predicted mean is (almost) equal
    - check if predicted variance is (almost) equal.
      All possible variances over the inputs and outputs are calculated
      and equality is checked.
    - check if variances within model are consistent. Parts of the covariance
      matrices should overlap, and this is tested.
    """

    elbos = [m.elbo(data) for m in models]

    # Check equality of log likelihood
    assert_all_array_elements_almost_equal(elbos)

    # Predict: full_cov = True and full_output_cov = True
    means_tt, vars_tt = predict_all(models, Data.Xs, full_cov=True, full_output_cov=True)
    # Predict: full_cov = True and full_output_cov = False
    means_tf, vars_tf = predict_all(models, Data.Xs, full_cov=True, full_output_cov=False)
    # Predict: full_cov = False and full_output_cov = True
    means_ft, vars_ft = predict_all(models, Data.Xs, full_cov=False, full_output_cov=True)
    # Predict: full_cov = False and full_output_cov = False
    means_ff, vars_ff = predict_all(models, Data.Xs, full_cov=False, full_output_cov=False)

    # check equality of all the means
    all_means = means_tt + means_tf + means_ft + means_ff
    assert_all_array_elements_almost_equal(all_means)

    # check equality of all the variances within a category
    # (e.g. full_cov=True and full_output_cov=False)
    for var in [vars_tt, vars_tf, vars_ft, vars_ff]:
        assert_all_array_elements_almost_equal(var)

    # Here we check that the variance in different categories are equal
    # after transforming to the right shape.
    var_tt = vars_tt[0]  # N x P x N x P
    var_tf = vars_tf[0]  # P x N x c
    var_ft = vars_ft[0]  # N x P x P
    var_ff = vars_ff[0]  # N x P

    np.testing.assert_almost_equal(
        np.diagonal(var_tt, axis1=1, axis2=3),
        np.transpose(var_tf, [1, 2, 0]),
        decimal=decimal,
    )
    np.testing.assert_almost_equal(
        np.diagonal(var_tt, axis1=0, axis2=2),
        np.transpose(var_ft, [1, 2, 0]),
        decimal=decimal,
    )
    np.testing.assert_almost_equal(
        np.diagonal(np.diagonal(var_tt, axis1=0, axis2=2)), var_ff, decimal=decimal
    )


@check_shapes(
    "q_sqrt: [L, M, M]",
    "W: [L, L]",
    "return: [1, LM, LM]",
)
def expand_cov(q_sqrt: tf.Tensor, W: tf.Tensor) -> tf.Tensor:
    """
    :param q_sqrt: cholesky of covariance matrices
    :param W: mixing matrix
    :return: cholesky of the expanded covariance matrix
    """
    q_cov = np.matmul(q_sqrt, q_sqrt.transpose([0, 2, 1]))  # [L, M, M]
    q_cov_expanded = scipy.linalg.block_diag(*q_cov)  # [LM, LM]
    q_sqrt_expanded = np.linalg.cholesky(q_cov_expanded)  # [LM, LM]
    return q_sqrt_expanded[None, ...]


@check_shapes(
    "return: [L, M, M]",
)
def create_q_sqrt(M: int, L: int) -> AnyNDArray:
    """ returns an array of lower triangular matrices """
    return np.array([np.tril(rng.randn(M, M)) for _ in range(L)])  # [L, M, M]


# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


class Data:
    cs = ShapeChecker().check_shape

    N, Ntest = 20, 5
    D = 1  # input dimension
    M = 3  # inducing points
    L = 2  # latent gps
    P = 3  # output dimension
    MAXITER = int(15e2)

    X = cs(tf.random.normal((N,), dtype=tf.float64)[:, None] * 10 - 5, "[N, 1]")
    G = cs(tf.concat([0.5 * tf.sin(3 * X) + X, 3.0 * tf.cos(X) - X], axis=1), "[N, L]")
    Ptrue = cs(tf.constant([[0.5, -0.3, 1.5], [-0.4, 0.43, 0.0]], dtype=tf.float64), "[L, P]")
    Y = cs(G @ Ptrue, "[N, P]")
    Y += cs(tf.random.normal(Y.shape, dtype=tf.float64) * [0.2, 0.2, 0.2], "[N, P]")
    Xs = cs(tf.linspace(-6, 6, Ntest)[:, None], "[Ntest, 1]")
    data = (X, Y)


class DataMixedKernelWithEye(Data):
    """ Note in this class L == P """

    cs = ShapeChecker().check_shape

    M, L = 4, 3
    W = cs(tf.eye(L, dtype=tf.float64), "[L, L]")

    G = cs(
        tf.concat(
            [0.5 * tf.sin(3 * Data.X) + Data.X, 3.0 * tf.cos(Data.X) - Data.X, 1.0 + Data.X], axis=1
        ),
        "[N, P]",
    )

    mu_data = cs(tf.random.uniform((M, L), dtype=tf.float64), "[M, L]")
    sqrt_data = cs(tf.convert_to_tensor(create_q_sqrt(M, L)), "[L, M, M]")

    mu_data_full = cs(tf.reshape(mu_data @ W, [-1, 1]), "[LM, 1]")
    sqrt_data_full = cs(
        tf.convert_to_tensor(expand_cov(sqrt_data.numpy(), W.numpy())), "[1, LM, LM]"
    )

    Y = cs(G @ W, "[N, L]")
    Y += cs(
        tf.random.normal(Y.shape, dtype=tf.float64) * tf.ones((L,), dtype=tf.float64) * 0.2,
        "[N, L]",
    )
    data = (Data.X, Y)


class DataMixedKernel(Data):
    cs = ShapeChecker().check_shape

    M = 5
    L = 2
    P = 3
    W = cs(tf.convert_to_tensor(rng.randn(P, L)), "[P, L]")
    G = cs(
        tf.concat([0.5 * tf.sin(3 * Data.X) + Data.X, 3.0 * tf.cos(Data.X) - Data.X], axis=1),
        "[N, L]",
    )

    mu_data = cs(tf.random.normal((M, L), dtype=tf.float64), "[M, L]")
    sqrt_data = cs(create_q_sqrt(M, L), "[L, M, M]")

    Y = cs(G @ tf.transpose(W), "[N, P]")
    Y += cs(
        tf.random.normal(Y.shape, dtype=tf.float64) * tf.ones((P,), dtype=tf.float64) * 0.1,
        "[N, P]",
    )
    data = (Data.X, Y)


# ------------------------------------------
# Test sample conditional
# ------------------------------------------


def test_sample_mvn(full_cov: bool) -> None:
    """
    Draws 10,000 samples from a distribution
    with known mean and covariance. The test checks
    if the mean and covariance of the samples is
    close to the true mean and covariance.
    """
    N, D = 10000, 2
    means = tf.ones((N, D), dtype=float_type)
    if full_cov:
        covs = tf.eye(D, batch_shape=[N], dtype=float_type)
    else:
        covs = tf.ones((N, D), dtype=float_type)

    samples = sample_mvn(means, covs, full_cov)
    samples_mean = np.mean(samples, axis=0)
    samples_cov = np.cov(samples, rowvar=False)

    np.testing.assert_array_almost_equal(samples_mean, [1.0, 1.0], decimal=1)
    np.testing.assert_array_almost_equal(samples_cov, [[1.0, 0.0], [0.0, 1.0]], decimal=1)


def test_sample_conditional(whiten: bool, full_cov: bool, full_output_cov: bool) -> None:
    if full_cov and full_output_cov:
        return

    q_mu = tf.random.uniform((Data.M, Data.P), dtype=tf.float64)  # [M, P]
    q_sqrt = tf.convert_to_tensor(
        [np.tril(tf.random.uniform((Data.M, Data.M), dtype=tf.float64)) for _ in range(Data.P)]
    )  # [P, M, M]

    Z = Data.X[: Data.M, ...]  # [M, D]
    Xs: AnyNDArray = np.ones((Data.N, Data.D), dtype=float_type)

    inducing_variable = InducingPoints(Z)
    kernel = SquaredExponential()

    # Path 1
    value_f, mean_f, var_f = sample_conditional(
        Xs,
        inducing_variable,
        kernel,
        q_mu,
        q_sqrt=q_sqrt,
        white=whiten,
        full_cov=full_cov,
        full_output_cov=full_output_cov,
        num_samples=int(1e5),
    )
    value_f = value_f.numpy().reshape((-1,) + value_f.numpy().shape[2:])

    # Path 2
    if full_output_cov:
        pytest.skip(
            "sample_conditional with X instead of inducing_variable does not support full_output_cov"
        )

    value_x, mean_x, var_x = sample_conditional(
        Xs,
        Z,
        kernel,
        q_mu,
        q_sqrt=q_sqrt,
        white=whiten,
        full_cov=full_cov,
        full_output_cov=full_output_cov,
        num_samples=int(1e5),
    )
    value_x = value_x.numpy().reshape((-1,) + value_x.numpy().shape[2:])

    # check if mean and covariance of samples are similar
    np.testing.assert_array_almost_equal(
        np.mean(value_x, axis=0), np.mean(value_f, axis=0), decimal=1
    )
    np.testing.assert_array_almost_equal(
        np.cov(value_x, rowvar=False), np.cov(value_f, rowvar=False), decimal=1
    )
    np.testing.assert_allclose(mean_x, mean_f)
    np.testing.assert_allclose(var_x, var_f)


def test_sample_conditional_mixedkernel() -> None:
    q_mu = tf.random.uniform((Data.M, Data.L), dtype=tf.float64)  # M x L
    q_sqrt = tf.convert_to_tensor(
        [np.tril(tf.random.uniform((Data.M, Data.M), dtype=tf.float64)) for _ in range(Data.L)]
    )  # L x M x M

    Z = Data.X[: Data.M, ...]  # M x D
    N = int(10e5)
    Xs: AnyNDArray = np.ones((N, Data.D), dtype=float_type)

    # Path 1: mixed kernel: most efficient route
    W = np.random.randn(Data.P, Data.L)
    mixed_kernel = mk.LinearCoregionalization([SquaredExponential() for _ in range(Data.L)], W)
    optimal_inducing_variable = mf.SharedIndependentInducingVariables(InducingPoints(Z))

    value, mean, var = sample_conditional(
        Xs, optimal_inducing_variable, mixed_kernel, q_mu, q_sqrt=q_sqrt, white=True
    )

    # Path 2: independent kernels, mixed later
    separate_kernel = mk.SeparateIndependent([SquaredExponential() for _ in range(Data.L)])
    fallback_inducing_variable = mf.SharedIndependentInducingVariables(InducingPoints(Z))

    value2, mean2, var2 = sample_conditional(
        Xs, fallback_inducing_variable, separate_kernel, q_mu, q_sqrt=q_sqrt, white=True
    )
    value2 = np.matmul(value2, W.T)
    # check if mean and covariance of samples are similar
    np.testing.assert_array_almost_equal(np.mean(value, axis=0), np.mean(value2, axis=0), decimal=1)
    np.testing.assert_array_almost_equal(
        np.cov(value, rowvar=False), np.cov(value2, rowvar=False), decimal=1
    )


QSqrtFactory = Callable[[tf.Tensor, int], tf.Tensor]


@pytest.fixture(
    name="fully_correlated_q_sqrt_factory",
    params=[lambda _, __: None, lambda LM, R: tf.eye(LM, batch_shape=(R,))],
)
def _q_sqrt_factory_fixture(request: SubRequest) -> QSqrtFactory:
    return cast(QSqrtFactory, request.param)


@pytest.mark.parametrize("R", [1, 2, 5])
def test_fully_correlated_conditional_repeat_shapes_fc_and_foc(
    R: int,
    fully_correlated_q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
) -> None:

    L, M, N, P = Data.L, Data.M, Data.N, Data.P

    Kmm = tf.ones((L * M, L * M)) + default_jitter() * tf.eye(L * M)
    Kmn = tf.ones((L * M, N, P))

    if full_cov and full_output_cov:
        Knn = tf.ones((N, P, N, P))
        expected_v_shape = [R, N, P, N, P]
    elif not full_cov and full_output_cov:
        Knn = tf.ones((N, P, P))
        expected_v_shape = [R, N, P, P]
    elif full_cov and not full_output_cov:
        Knn = tf.ones((P, N, N))
        expected_v_shape = [R, P, N, N]
    else:
        Knn = tf.ones((N, P))
        expected_v_shape = [R, N, P]

    f = tf.ones((L * M, R))
    q_sqrt = fully_correlated_q_sqrt_factory(L * M, R)

    m, v = fully_correlated_conditional_repeat(
        Kmn,
        Kmm,
        Knn,
        f,
        full_cov=full_cov,
        full_output_cov=full_output_cov,
        q_sqrt=q_sqrt,
        white=whiten,
    )

    assert m.shape.as_list() == [R, N, P]
    assert v.shape.as_list() == expected_v_shape


def test_fully_correlated_conditional_repeat_whiten(whiten: bool) -> None:
    """
    This test checks the effect of the `white` flag, which changes the projection matrix `A`.

    The impact of the flag on the value of `A` can be easily verified by its effect on the
    predicted mean. While the predicted covariance is also a function of `A` this test does not
    inspect that value.
    """
    N, P = Data.N, Data.P

    Lm: AnyNDArray = np.random.randn(1, 1).astype(np.float32) ** 2
    Kmm = Lm * Lm + default_jitter()

    Kmn = tf.ones((1, N, P))

    Knn = tf.ones((N, P))
    f: AnyNDArray = np.random.randn(1, 1).astype(np.float32)

    mean, _ = fully_correlated_conditional_repeat(
        Kmn,
        Kmm,
        Knn,
        f,
        white=whiten,
    )

    if whiten:
        expected_mean = (f * Kmn) / Lm
    else:
        expected_mean = (f * Kmn) / Kmm

    np.testing.assert_allclose(mean, expected_mean, rtol=1e-3)


def test_fully_correlated_conditional_shapes_fc_and_foc(
    fully_correlated_q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
) -> None:
    L, M, N, P = Data.L, Data.M, Data.N, Data.P

    Kmm = tf.ones((L * M, L * M)) + default_jitter() * tf.eye(L * M)
    Kmn = tf.ones((L * M, N, P))

    if full_cov and full_output_cov:
        Knn = tf.ones((N, P, N, P))
        expected_v_shape = [N, P, N, P]
    elif not full_cov and full_output_cov:
        Knn = tf.ones((N, P, P))
        expected_v_shape = [N, P, P]
    elif full_cov and not full_output_cov:
        Knn = tf.ones((P, N, N))
        expected_v_shape = [P, N, N]
    else:
        Knn = tf.ones((N, P))
        expected_v_shape = [N, P]

    f = tf.ones((L * M, 1))
    q_sqrt = fully_correlated_q_sqrt_factory(L * M, 1)

    m, v = fully_correlated_conditional(
        Kmn,
        Kmm,
        Knn,
        f,
        full_cov=full_cov,
        full_output_cov=full_output_cov,
        q_sqrt=q_sqrt,
        white=whiten,
    )

    assert m.shape.as_list() == [N, P]
    assert v.shape.as_list() == expected_v_shape


# ------------------------------------------
# Test Mok Output Dims
# ------------------------------------------


def test_shapes_of_mok() -> None:
    data = DataMixedKernel

    kern_list = [SquaredExponential() for _ in range(data.L)]

    k1 = mk.LinearCoregionalization(kern_list, W=data.W)
    assert k1.num_latent_gps == data.L

    k2 = mk.SeparateIndependent(kern_list)
    assert k2.num_latent_gps == data.L

    dims = 5
    k3 = mk.SharedIndependent(SquaredExponential(), dims)
    assert k3.num_latent_gps == dims


# ------------------------------------------
# Test Mixed Mok Kgg
# ------------------------------------------


def test_MixedMok_Kgg() -> None:
    data = DataMixedKernel
    kern_list = [SquaredExponential() for _ in range(data.L)]
    kernel = mk.LinearCoregionalization(kern_list, W=data.W)

    Kgg = kernel.Kgg(Data.X, Data.X)  # L x N x N
    Kff = kernel.K(Data.X, Data.X)  # N x P x N x P

    # Kff = W @ Kgg @ W^T
    Kff_infered = np.einsum("lnm,pl,ql->npmq", Kgg, data.W, data.W)

    np.testing.assert_array_almost_equal(Kff, Kff_infered, decimal=5)


# ------------------------------------------
# Integration tests
# ------------------------------------------


def test_shared_independent_mok() -> None:
    """
    In this test we use the same kernel and the same inducing inducing
    for each of the outputs. The outputs are considered to be uncorrelated.
    This is how GPflow handled multiple outputs before the multioutput framework was added.
    We compare three models here:
        1) an ineffient one, where we use a SharedIndepedentMok with InducingPoints.
           This combination will uses a Kff of size N x P x N x P, Kfu if size N x P x M x P
           which is extremely inefficient as most of the elements are zero.
        2) efficient: SharedIndependentMok and SharedIndependentMof
           This combinations uses the most efficient form of matrices
        3) the old way, efficient way: using Kernel and InducingPoints
        Model 2) and 3) follow more or less the same code path.
    """
    np.random.seed(0)
    # Model 1
    q_mu_1 = np.random.randn(Data.M * Data.P, 1)  # MP x 1
    q_sqrt_1 = np.tril(np.random.randn(Data.M * Data.P, Data.M * Data.P))[None, ...]  # 1 x MP x MP
    kernel_1 = mk.SharedIndependent(SquaredExponential(variance=0.5, lengthscales=1.2), Data.P)
    inducing_variable = InducingPoints(Data.X[: Data.M, ...])
    model_1 = SVGP(
        kernel_1,
        Gaussian(),
        inducing_variable,
        q_mu=q_mu_1,
        q_sqrt=q_sqrt_1,
        num_latent_gps=Data.Y.shape[-1],
    )
    set_trainable(model_1, False)
    set_trainable(model_1.q_sqrt, True)

    gpflow.optimizers.Scipy().minimize(
        model_1.training_loss_closure(Data.data),
        variables=model_1.trainable_variables,
        options=dict(maxiter=500),
        method="BFGS",
        compile=True,
    )

    # Model 2
    q_mu_2 = np.reshape(q_mu_1, [Data.M, Data.P])  # M x P
    q_sqrt_2: AnyNDArray = np.array(
        [np.tril(np.random.randn(Data.M, Data.M)) for _ in range(Data.P)]
    )  # P x M x M
    kernel_2 = SquaredExponential(variance=0.5, lengthscales=1.2)
    inducing_variable_2 = InducingPoints(Data.X[: Data.M, ...])
    model_2 = SVGP(
        kernel_2,
        Gaussian(),
        inducing_variable_2,
        num_latent_gps=Data.P,
        q_mu=q_mu_2,
        q_sqrt=q_sqrt_2,
    )
    set_trainable(model_2, False)
    set_trainable(model_2.q_sqrt, True)

    gpflow.optimizers.Scipy().minimize(
        model_2.training_loss_closure(Data.data),
        variables=model_2.trainable_variables,
        options=dict(maxiter=500),
        method="BFGS",
        compile=True,
    )

    # Model 3
    q_mu_3 = np.reshape(q_mu_1, [Data.M, Data.P])  # M x P
    q_sqrt_3: AnyNDArray = np.array(
        [np.tril(np.random.randn(Data.M, Data.M)) for _ in range(Data.P)]
    )  # P x M x M
    kernel_3 = mk.SharedIndependent(SquaredExponential(variance=0.5, lengthscales=1.2), Data.P)
    inducing_variable_3 = mf.SharedIndependentInducingVariables(
        InducingPoints(Data.X[: Data.M, ...])
    )
    model_3 = SVGP(
        kernel_3,
        Gaussian(),
        inducing_variable_3,
        num_latent_gps=Data.P,
        q_mu=q_mu_3,
        q_sqrt=q_sqrt_3,
    )
    set_trainable(model_3, False)
    set_trainable(model_3.q_sqrt, True)

    gpflow.optimizers.Scipy().minimize(
        model_3.training_loss_closure(Data.data),
        variables=model_3.trainable_variables,
        options=dict(maxiter=500),
        method="BFGS",
        compile=True,
    )

    check_equality_predictions(Data.data, [model_1, model_2, model_3])


def test_separate_independent_mok() -> None:
    """
    We use different independent kernels for each of the output dimensions.
    We can achieve this in two ways:
        1) efficient: SeparateIndependentMok with Shared/SeparateIndependentMof
        2) inefficient: SeparateIndependentMok with InducingPoints
    However, both methods should return the same conditional,
    and after optimization return the same log likelihood.
    """
    # Model 1 (Inefficient)
    q_mu_1 = np.random.randn(Data.M * Data.P, 1)
    q_sqrt_1 = np.tril(np.random.randn(Data.M * Data.P, Data.M * Data.P))[None, ...]  # 1 x MP x MP

    kern_list_1 = [SquaredExponential(variance=0.5, lengthscales=1.2) for _ in range(Data.P)]
    kernel_1 = mk.SeparateIndependent(kern_list_1)
    inducing_variable_1 = InducingPoints(Data.X[: Data.M, ...])
    model_1 = SVGP(
        kernel_1,
        Gaussian(),
        inducing_variable_1,
        num_latent_gps=1,
        q_mu=q_mu_1,
        q_sqrt=q_sqrt_1,
    )
    set_trainable(model_1, False)
    set_trainable(model_1.q_sqrt, True)
    set_trainable(model_1.q_mu, True)

    gpflow.optimizers.Scipy().minimize(
        model_1.training_loss_closure(Data.data),
        variables=model_1.trainable_variables,
        method="BFGS",
        compile=True,
    )

    # Model 2 (efficient)
    q_mu_2 = np.random.randn(Data.M, Data.P)
    q_sqrt_2: AnyNDArray = np.array(
        [np.tril(np.random.randn(Data.M, Data.M)) for _ in range(Data.P)]
    )  # P x M x M
    kern_list_2 = [SquaredExponential(variance=0.5, lengthscales=1.2) for _ in range(Data.P)]
    kernel_2 = mk.SeparateIndependent(kern_list_2)
    inducing_variable_2 = mf.SharedIndependentInducingVariables(
        InducingPoints(Data.X[: Data.M, ...])
    )
    model_2 = SVGP(
        kernel_2,
        Gaussian(),
        inducing_variable_2,
        num_latent_gps=Data.P,
        q_mu=q_mu_2,
        q_sqrt=q_sqrt_2,
    )
    set_trainable(model_2, False)
    set_trainable(model_2.q_sqrt, True)
    set_trainable(model_2.q_mu, True)

    gpflow.optimizers.Scipy().minimize(
        model_2.training_loss_closure(Data.data),
        variables=model_2.trainable_variables,
        method="BFGS",
        compile=True,
    )

    check_equality_predictions(Data.data, [model_1, model_2])


def test_separate_independent_mof() -> None:
    """
    Same test as above but we use different (i.e. separate) inducing inducing
    for each of the output dimensions.
    """
    np.random.seed(0)

    # Model 1 (INefficient)
    q_mu_1 = np.random.randn(Data.M * Data.P, 1)
    q_sqrt_1 = np.tril(np.random.randn(Data.M * Data.P, Data.M * Data.P))[None, ...]  # 1 x MP x MP

    kernel_1 = mk.SharedIndependent(SquaredExponential(variance=0.5, lengthscales=1.2), Data.P)
    inducing_variable_1 = InducingPoints(Data.X[: Data.M, ...])
    model_1 = SVGP(kernel_1, Gaussian(), inducing_variable_1, q_mu=q_mu_1, q_sqrt=q_sqrt_1)
    set_trainable(model_1, False)
    set_trainable(model_1.q_sqrt, True)
    set_trainable(model_1.q_mu, True)

    gpflow.optimizers.Scipy().minimize(
        model_1.training_loss_closure(Data.data),
        variables=model_1.trainable_variables,
        method="BFGS",
        compile=True,
    )

    # Model 2 (efficient)
    q_mu_2 = np.random.randn(Data.M, Data.P)
    q_sqrt_2: AnyNDArray = np.array(
        [np.tril(np.random.randn(Data.M, Data.M)) for _ in range(Data.P)]
    )  # P x M x M
    kernel_2 = mk.SharedIndependent(SquaredExponential(variance=0.5, lengthscales=1.2), Data.P)
    inducing_variable_list_2 = [InducingPoints(Data.X[: Data.M, ...]) for _ in range(Data.P)]
    inducing_variable_2 = mf.SeparateIndependentInducingVariables(inducing_variable_list_2)
    model_2 = SVGP(kernel_2, Gaussian(), inducing_variable_2, q_mu=q_mu_2, q_sqrt=q_sqrt_2)
    set_trainable(model_2, False)
    set_trainable(model_2.q_sqrt, True)
    set_trainable(model_2.q_mu, True)

    gpflow.optimizers.Scipy().minimize(
        model_2.training_loss_closure(Data.data),
        variables=model_2.trainable_variables,
        method="BFGS",
        compile=True,
    )

    # Model 3 (Inefficient): an idenitical inducing variable is used P times,
    # and treated as a separate one.
    q_mu_3 = np.random.randn(Data.M, Data.P)
    q_sqrt_3: AnyNDArray = np.array(
        [np.tril(np.random.randn(Data.M, Data.M)) for _ in range(Data.P)]
    )  # P x M x M
    kern_list = [SquaredExponential(variance=0.5, lengthscales=1.2) for _ in range(Data.P)]
    kernel_3 = mk.SeparateIndependent(kern_list)
    inducing_variable_list_3 = [InducingPoints(Data.X[: Data.M, ...]) for _ in range(Data.P)]
    inducing_variable_3 = mf.SeparateIndependentInducingVariables(inducing_variable_list_3)
    model_3 = SVGP(kernel_3, Gaussian(), inducing_variable_3, q_mu=q_mu_3, q_sqrt=q_sqrt_3)
    set_trainable(model_3, False)
    set_trainable(model_3.q_sqrt, True)
    set_trainable(model_3.q_mu, True)

    gpflow.optimizers.Scipy().minimize(
        model_3.training_loss_closure(Data.data),
        variables=model_3.trainable_variables,
        method="BFGS",
        compile=True,
    )

    check_equality_predictions(Data.data, [model_1, model_2, model_3])


def test_mixed_mok_with_Id_vs_independent_mok() -> None:
    data = DataMixedKernelWithEye
    # Independent model
    k1 = mk.SharedIndependent(SquaredExponential(variance=0.5, lengthscales=1.2), data.L)
    f1 = InducingPoints(data.X[: data.M, ...])
    model_1 = SVGP(k1, Gaussian(), f1, q_mu=data.mu_data_full, q_sqrt=data.sqrt_data_full)
    set_trainable(model_1, False)
    set_trainable(model_1.q_sqrt, True)

    gpflow.optimizers.Scipy().minimize(
        model_1.training_loss_closure(Data.data),
        variables=model_1.trainable_variables,
        method="BFGS",
        compile=True,
    )

    # Mixed Model
    kern_list = [SquaredExponential(variance=0.5, lengthscales=1.2) for _ in range(data.L)]
    k2 = mk.LinearCoregionalization(kern_list, data.W)
    f2 = InducingPoints(data.X[: data.M, ...])
    model_2 = SVGP(k2, Gaussian(), f2, q_mu=data.mu_data_full, q_sqrt=data.sqrt_data_full)
    set_trainable(model_2, False)
    set_trainable(model_2.q_sqrt, True)

    gpflow.optimizers.Scipy().minimize(
        model_2.training_loss_closure(Data.data),
        variables=model_2.trainable_variables,
        method="BFGS",
        compile=True,
    )

    check_equality_predictions(Data.data, [model_1, model_2])


def test_compare_mixed_kernel() -> None:
    data = DataMixedKernel

    kern_list = [SquaredExponential() for _ in range(data.L)]
    k1 = mk.LinearCoregionalization(kern_list, W=data.W)
    f1 = mf.SharedIndependentInducingVariables(InducingPoints(data.X[: data.M, ...]))
    model_1 = SVGP(k1, Gaussian(), inducing_variable=f1, q_mu=data.mu_data, q_sqrt=data.sqrt_data)

    kern_list = [SquaredExponential() for _ in range(data.L)]
    k2 = mk.LinearCoregionalization(kern_list, W=data.W)
    f2 = mf.SharedIndependentInducingVariables(InducingPoints(data.X[: data.M, ...]))
    model_2 = SVGP(k2, Gaussian(), inducing_variable=f2, q_mu=data.mu_data, q_sqrt=data.sqrt_data)

    check_equality_predictions(Data.data, [model_1, model_2])


def test_multioutput_with_diag_q_sqrt() -> None:
    data = DataMixedKernel

    q_sqrt_diag: AnyNDArray = np.ones((data.M, data.L)) * 2
    q_sqrt: AnyNDArray = np.repeat(np.eye(data.M)[None, ...], data.L, axis=0) * 2  # L x M x M

    kern_list = [SquaredExponential() for _ in range(data.L)]
    k1 = mk.LinearCoregionalization(kern_list, W=data.W)
    f1 = mf.SharedIndependentInducingVariables(InducingPoints(data.X[: data.M, ...]))
    model_1 = SVGP(
        k1,
        Gaussian(),
        inducing_variable=f1,
        q_mu=data.mu_data,
        q_sqrt=q_sqrt_diag,
        q_diag=True,
    )

    kern_list = [SquaredExponential() for _ in range(data.L)]
    k2 = mk.LinearCoregionalization(kern_list, W=data.W)
    f2 = mf.SharedIndependentInducingVariables(InducingPoints(data.X[: data.M, ...]))
    model_2 = SVGP(
        k2,
        Gaussian(),
        inducing_variable=f2,
        q_mu=data.mu_data,
        q_sqrt=q_sqrt,
        q_diag=False,
    )

    check_equality_predictions(Data.data, [model_1, model_2])


def test_MixedKernelSeparateMof() -> None:
    data = DataMixedKernel

    kern_list = [SquaredExponential() for _ in range(data.L)]
    inducing_variable_list = [InducingPoints(data.X[: data.M, ...]) for _ in range(data.L)]
    k1 = mk.LinearCoregionalization(kern_list, W=data.W)
    f1 = mf.SeparateIndependentInducingVariables(inducing_variable_list)
    model_1 = SVGP(k1, Gaussian(), inducing_variable=f1, q_mu=data.mu_data, q_sqrt=data.sqrt_data)

    kern_list = [SquaredExponential() for _ in range(data.L)]
    inducing_variable_list = [InducingPoints(data.X[: data.M, ...]) for _ in range(data.L)]
    k2 = mk.LinearCoregionalization(kern_list, W=data.W)
    f2 = mf.SeparateIndependentInducingVariables(inducing_variable_list)
    model_2 = SVGP(k2, Gaussian(), inducing_variable=f2, q_mu=data.mu_data, q_sqrt=data.sqrt_data)

    check_equality_predictions(Data.data, [model_1, model_2])


def test_separate_independent_conditional_with_q_sqrt_none() -> None:
    """
    In response to bug #1523, this test checks that separate_independent_condtional
    does not fail when q_sqrt=None.
    """
    q_sqrt = None
    data = DataMixedKernel

    kern_list = [SquaredExponential() for _ in range(data.L)]
    kernel = gpflow.kernels.SeparateIndependent(kern_list)
    inducing_variable_list = [InducingPoints(data.X[: data.M, ...]) for _ in range(data.L)]
    inducing_variable = mf.SeparateIndependentInducingVariables(inducing_variable_list)

    gpflow.conditionals.conditional(
        data.X,
        inducing_variable,
        kernel,
        data.mu_data,
        full_cov=False,
        full_output_cov=False,
        q_sqrt=q_sqrt,
        white=True,
    )


def test_independent_interdomain_conditional_bug_regression() -> None:
    """
    Regression test for https://github.com/GPflow/GPflow/issues/818
    Not an exhaustive test
    """
    cs = ShapeChecker().check_shape

    M = 31
    N = 11
    D_lat = 5
    D_inp = D_lat * 7
    L = 2
    P = 3

    X = cs(np.random.randn(N, D_inp), "[N, D_inp]")
    Zs = cs([np.random.randn(M, D_lat) for _ in range(L)], "[L, M, D_lat]")
    k = gpflow.kernels.SquaredExponential(lengthscales=np.ones(D_lat))

    @check_shapes(
        "Z: [M, D_lat]",
        "X: [N, D_inp]",
        "return: [P, M, N]",
    )
    def compute_Kmn(Z: tf.Tensor, X: tf.Tensor) -> tf.Tensor:
        return tf.stack([k(Z, X[:, i * D_lat : (i + 1) * D_lat]) for i in range(P)])

    @check_shapes(
        "X: [N, D_inp]",
        "return: [P, N]",
    )
    def compute_Knn(X: tf.Tensor) -> tf.Tensor:
        return tf.stack([k(X[:, i * D_lat : (i + 1) * D_lat], full_cov=False) for i in range(P)])

    Kmm = cs(tf.stack([k(Z) for Z in Zs]), "[L, M, M]")
    Kmn = cs(tf.stack([compute_Kmn(Z, X) for Z in Zs]), "[L, P, M, N]")
    Kmn = cs(tf.transpose(Kmn, [2, 0, 3, 1]), "[M, L, N, P]")
    Knn = cs(tf.transpose(compute_Knn(X)), "[N, P]")
    q_mu = cs(tf.convert_to_tensor(np.zeros((M, L))), "[M, L]")
    q_sqrt = cs(tf.convert_to_tensor(np.stack([np.eye(M) for _ in range(L)])), "[L, M, M]")

    _, _ = independent_interdomain_conditional(
        Kmn, Kmm, Knn, q_mu, q_sqrt=q_sqrt, full_cov=False, full_output_cov=False
    )


def test_independent_interdomain_conditional_whiten(whiten: bool) -> None:
    """
    This test checks the effect of the `white` flag, which changes the projection matrix `A`.

    The impact of the flag on the value of `A` can be easily verified by its effect on the
    predicted mean. While the predicted covariance is also a function of `A` this test does not
    inspect that value.
    """
    N, P = Data.N, Data.P

    Lm: AnyNDArray = np.random.randn(1, 1, 1).astype(np.float32) ** 2
    Kmm = Lm * Lm + default_jitter()

    Kmn = tf.ones((1, 1, N, P))

    Knn = tf.ones((N, P))
    f: AnyNDArray = np.random.randn(1, 1).astype(np.float32)

    mean, _ = independent_interdomain_conditional(
        Kmn,
        Kmm,
        Knn,
        f,
        white=whiten,
    )

    if whiten:
        expected_mean = (f * Kmn) / Lm
    else:
        expected_mean = (f * Kmn) / Kmm

    np.testing.assert_allclose(mean, expected_mean[0][0], rtol=1e-2)


# tests/gpflow/conditionals/test_uncertain_conditional.py

# Copyright 2017 the GPflow authors.
#
# 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 typing import Collection, Optional

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
from gpflow.base import AnyNDArray, MeanAndVariance
from gpflow.conditionals import conditional, uncertain_conditional
from gpflow.config import default_float
from gpflow.experimental.check_shapes import ShapeChecker, check_shapes
from gpflow.mean_functions import Constant, Linear, MeanFunction, Zero
from gpflow.quadrature import mvnquad
from gpflow.utilities import training_loop

rng = np.random.RandomState(1)

# ------------------------------------------
# Helpers
# ------------------------------------------


class MomentMatchingSVGP(gpflow.models.SVGP):
    @check_shapes(
        "Xmu: [batch..., N, Din]",
        "Xcov: [batch..., N, n, n]",
        "self.inducing_variable: [M, Din, broadcast t]",
        "self.q_mu: [M, Dout]",
        "self.q_sqrt: [t, M, M]",
        "return[0]: [batch..., N, Dout]",
        "return[1]: [batch..., N, t, t] if self.full_output_cov",
        "return[1]: [batch..., N, Dout] if not self.full_output_cov",
    )
    def uncertain_predict_f_moment_matching(
        self, Xmu: tf.Tensor, Xcov: tf.Tensor
    ) -> MeanAndVariance:
        return uncertain_conditional(
            Xmu,
            Xcov,
            self.inducing_variable,
            self.kernel,
            self.q_mu,
            self.q_sqrt,
            mean_function=self.mean_function,
            white=self.whiten,
            full_output_cov=self.full_output_cov,
        )

    @check_shapes(
        "Xmu: [Din]",
        "Xchol: [n, n]",
        "return[0]: [Dout]",
        "return[1]: [t, t] if self.full_output_cov",
        "return[1]: [Dout] if not self.full_output_cov",
    )
    def uncertain_predict_f_monte_carlo(
        self, Xmu: tf.Tensor, Xchol: tf.Tensor, mc_iter: int = int(1e6)
    ) -> MeanAndVariance:
        D_in = Xchol.shape[0]
        X_samples = Xmu + np.reshape(
            Xchol[None, :, :] @ rng.randn(mc_iter, D_in)[:, :, None], [mc_iter, D_in]
        )
        F_mu, F_var = self.predict_f(X_samples)
        F_samples = (F_mu + rng.randn(*F_var.shape) * (F_var ** 0.5)).numpy()
        mean = np.mean(F_samples, axis=0)
        covar = np.cov(F_samples.T)
        return mean, covar


@check_shapes(
    "return: [n, shape...]",
)
def gen_L(n: int, *shape: int) -> AnyNDArray:
    return np.array([np.tril(rng.randn(*shape)) for _ in range(n)])


def gen_q_sqrt(D_out: int, *shape: int) -> tf.Tensor:
    return tf.convert_to_tensor(
        np.array([np.tril(rng.randn(*shape)) for _ in range(D_out)]),
        dtype=default_float(),
    )


def mean_function_factory(
    mean_function_name: Optional[str], D_in: int, D_out: int
) -> Optional[MeanFunction]:
    if mean_function_name == "Zero":
        return Zero(output_dim=D_out)
    elif mean_function_name == "Constant":
        return Constant(c=rng.rand(D_out))
    elif mean_function_name == "Linear":
        return Linear(A=rng.rand(D_in, D_out), b=rng.rand(D_out))
    else:
        return None


# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


class Data:
    cs = ShapeChecker().check_shape

    N = 7
    N_new = 2
    D_out = 3
    D_in = 1
    X = cs(np.linspace(-5, 5, N)[:, None] + rng.randn(N, 1), "[N, D_in]")
    Y: AnyNDArray = cs(np.hstack([np.sin(X), np.cos(X), X ** 2]), "[N, D_out]")
    Xnew_mu = cs(rng.randn(N_new, 1), "[N_new, D_in]")
    Xnew_covar = cs(np.zeros((N_new, 1, 1)), "[N_new, D_in, D_in]")
    data = (X, Y)


class DataMC1(Data):
    Y: AnyNDArray = Data.cs(  # type: ignore[misc]
        np.hstack([np.sin(Data.X), np.sin(Data.X) * 2, Data.X ** 2]), "[N, D_out]"
    )
    data = (Data.X, Y)


class DataMC2(Data):
    cs = ShapeChecker().check_shape

    N = 7
    N_new = 5
    D_out = 4
    D_in = 2
    X = cs(rng.randn(N, D_in), "[N, D_in]")
    Y: AnyNDArray = cs(np.hstack([np.sin(X), np.sin(X)]), "[N, D_out]")
    Xnew_mu = cs(rng.randn(N_new, D_in), "[N_new, D_in]")
    L = cs(gen_L(N_new, D_in, D_in), "[N_new, D_in, D_in]")
    Xnew_covar: AnyNDArray = cs(np.array([l @ l.T for l in L]), "[N_new, D_in, D_in]")
    data = (X, Y)


class DataQuad:
    cs = ShapeChecker().check_shape

    num_data = 10
    num_ind = 10
    D_in = 2
    D_out = 3
    H = 150
    Xmu = cs(
        tf.convert_to_tensor(rng.randn(num_data, D_in), dtype=default_float()), "[num_data, D_in]"
    )
    L = cs(gen_L(num_data, D_in, D_in), "[num_data, D_in, D_in]")
    Xvar = cs(
        tf.convert_to_tensor(np.array([l @ l.T for l in L]), dtype=default_float()),
        "[num_data, D_in, D_in]",
    )
    Z = cs(rng.randn(num_ind, D_in), "[num_ind, D_in]")
    q_mu = cs(
        tf.convert_to_tensor(rng.randn(num_ind, D_out), dtype=default_float()), "[num_ind, D_out]"
    )
    q_sqrt = cs(gen_q_sqrt(D_out, num_ind, num_ind), "[D_out, num_ind, num_ind]")


MEANS: Collection[Optional[str]] = ["Constant", "Linear", "Zero", None]


@pytest.mark.parametrize("white", [True, False])
@pytest.mark.parametrize("mean", MEANS)
def test_no_uncertainty(white: bool, mean: Optional[str]) -> None:
    mean_function = mean_function_factory(mean, Data.D_in, Data.D_out)
    kernel = gpflow.kernels.SquaredExponential(variance=rng.rand())
    model = MomentMatchingSVGP(
        kernel,
        gpflow.likelihoods.Gaussian(),
        num_latent_gps=Data.D_out,
        mean_function=mean_function,
        inducing_variable=Data.X.copy(),
        whiten=white,
    )
    model.full_output_cov = False

    training_loop(
        model.training_loss_closure(Data.data),
        optimizer=tf.optimizers.Adam(),
        var_list=model.trainable_variables,
        maxiter=100,
        compile=True,
    )

    mean1, var1 = model.predict_f(Data.Xnew_mu)
    mean2, var2 = model.uncertain_predict_f_moment_matching(
        *map(tf.convert_to_tensor, [Data.Xnew_mu, Data.Xnew_covar])
    )

    assert_allclose(mean1, mean2)
    for n in range(Data.N_new):
        assert_allclose(var1[n, :], var2[n, ...])


@pytest.mark.parametrize("white", [True, False])
@pytest.mark.parametrize("mean", MEANS)
def test_monte_carlo_1_din(white: bool, mean: Optional[str]) -> None:
    kernel = gpflow.kernels.SquaredExponential(variance=rng.rand())
    mean_function = mean_function_factory(mean, DataMC1.D_in, DataMC1.D_out)
    model = MomentMatchingSVGP(
        kernel,
        gpflow.likelihoods.Gaussian(),
        num_latent_gps=DataMC1.D_out,
        mean_function=mean_function,
        inducing_variable=DataMC1.X.copy(),
        whiten=white,
    )
    model.full_output_cov = True

    training_loop(
        model.training_loss_closure(DataMC1.data),
        optimizer=tf.optimizers.Adam(),
        var_list=model.trainable_variables,
        maxiter=200,
        compile=True,
    )

    mean1, var1 = model.uncertain_predict_f_moment_matching(
        *map(tf.convert_to_tensor, [DataMC1.Xnew_mu, DataMC1.Xnew_covar])
    )

    for n in range(DataMC1.N_new):
        mean2, var2 = model.uncertain_predict_f_monte_carlo(
            DataMC1.Xnew_mu[n, ...], DataMC1.Xnew_covar[n, ...] ** 0.5
        )
        assert_allclose(mean1[n, ...], mean2, atol=1e-3, rtol=1e-1)
        assert_allclose(var1[n, ...], var2, atol=1e-2, rtol=1e-1)


@pytest.mark.parametrize("white", [True, False])
@pytest.mark.parametrize("mean", MEANS)
def test_monte_carlo_2_din(white: bool, mean: Optional[str]) -> None:
    kernel = gpflow.kernels.SquaredExponential(variance=rng.rand())
    mean_function = mean_function_factory(mean, DataMC2.D_in, DataMC2.D_out)
    model = MomentMatchingSVGP(
        kernel,
        gpflow.likelihoods.Gaussian(),
        num_latent_gps=DataMC2.D_out,
        mean_function=mean_function,
        inducing_variable=DataMC2.X.copy(),
        whiten=white,
    )
    model.full_output_cov = True

    training_loop(
        model.training_loss_closure(DataMC2.data),
        optimizer=tf.optimizers.Adam(),
        var_list=model.trainable_variables,
        maxiter=100,
        compile=True,
    )

    mean1, var1 = model.uncertain_predict_f_moment_matching(
        *map(tf.convert_to_tensor, [DataMC2.Xnew_mu, DataMC2.Xnew_covar])
    )

    for n in range(DataMC2.N_new):
        mean2, var2 = model.uncertain_predict_f_monte_carlo(
            DataMC2.Xnew_mu[n, ...], DataMC2.L[n, ...]
        )
        assert_allclose(mean1[n, ...], mean2, atol=1e-2)
        assert_allclose(var1[n, ...], var2, atol=1e-2)


@pytest.mark.parametrize("white", [True, False])
@pytest.mark.parametrize("mean", MEANS)
def test_quadrature(white: bool, mean: Optional[str]) -> None:
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = gpflow.inducing_variables.InducingPoints(DataQuad.Z)
    mean_function = mean_function_factory(mean, DataQuad.D_in, DataQuad.D_out)

    effective_mean = mean_function or (lambda X: 0.0)

    def conditional_fn(X: tf.Tensor) -> tf.Tensor:
        return conditional(
            X,
            inducing_variable,
            kernel,
            DataQuad.q_mu,
            q_sqrt=DataQuad.q_sqrt,
            white=white,
        )

    def mean_fn(X: tf.Tensor) -> tf.Tensor:
        return conditional_fn(X)[0] + effective_mean(X)

    def var_fn(X: tf.Tensor) -> tf.Tensor:
        return conditional_fn(X)[1]

    quad_args = (
        DataQuad.Xmu,
        DataQuad.Xvar,
        DataQuad.H,
        DataQuad.D_in,
        (DataQuad.D_out,),
    )
    mean_quad = mvnquad(mean_fn, *quad_args)
    var_quad = mvnquad(var_fn, *quad_args)

    def mean_sq_fn(X: tf.Tensor) -> tf.Tensor:
        return mean_fn(X) ** 2

    mean_sq_quad = mvnquad(mean_sq_fn, *quad_args)
    var_quad = var_quad + (mean_sq_quad - mean_quad ** 2)

    mean_analytic, var_analytic = uncertain_conditional(
        DataQuad.Xmu,
        DataQuad.Xvar,
        inducing_variable,
        kernel,
        DataQuad.q_mu,
        DataQuad.q_sqrt,
        mean_function=mean_function,
        full_output_cov=False,
        white=white,
    )

    assert_allclose(mean_quad, mean_analytic, rtol=1e-6)
    assert_allclose(var_quad, var_analytic, rtol=1e-6)


# tests/gpflow/conditionals/test_util.py

# Copyright 2020 the GPflow authors.
#
# 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 typing import Optional, Tuple

import numpy as np
import pytest
import tensorflow as tf
from _pytest.logging import LogCaptureFixture
from numpy.testing import assert_allclose, assert_equal

from gpflow import default_float
from gpflow.conditionals.util import rollaxis_left, rollaxis_right, sample_mvn
from gpflow.utilities.ops import leading_transpose


def test_leading_transpose() -> None:
    dims = [1, 2, 3, 4]
    a = tf.zeros(dims)
    b = leading_transpose(a, [..., -1, -2])
    c = leading_transpose(a, [-1, ..., -2])
    d = leading_transpose(a, [-1, -2, ...])
    e = leading_transpose(a, [3, 2, ...])
    f = leading_transpose(a, [3, -2, ...])

    assert len(a.shape) == len(b.shape) == len(c.shape) == len(d.shape)
    assert len(a.shape) == len(e.shape) == len(f.shape)
    assert b.shape[-2:] == [4, 3]
    assert c.shape[0] == 4 and c.shape[-1] == 3
    assert d.shape[:2] == [4, 3]
    assert d.shape == e.shape == f.shape


def test_leading_transpose_fails() -> None:
    """ Check that error is thrown if `perm` is not compatible with `a` """
    dims = [1, 2, 3, 4]
    a = tf.zeros(dims)

    with pytest.raises(ValueError):
        leading_transpose(a, [-1, -2])


def test_leading_transpose_with_tf_function_wrapper(caplog: LogCaptureFixture) -> None:
    """ Check that no warnings are thrown when compiling `leading_transpose` """

    dims = [1, 2, 3, 4]
    a = tf.zeros(dims)

    @tf.function
    def compiled_wrapper() -> tf.Tensor:
        return leading_transpose(a, [..., -1, -2])

    compiled_wrapper()

    # When Autograph cannot compile a function it sends a WARNING message to the
    # logs.
    log_message_levels = [record.levelname for record in caplog.records]
    assert "WARNING" not in log_message_levels


# rollaxis
@pytest.mark.parametrize("rolls", [1, 2])
@pytest.mark.parametrize("direction", ["left", "right"])
def test_rollaxis(rolls: int, direction: str) -> None:
    A = np.random.randn(10, 5, 3)
    A_tf = tf.convert_to_tensor(A)

    if direction == "left":
        perm = [1, 2, 0] if rolls == 1 else [2, 0, 1]
    elif direction == "right":
        perm = [2, 0, 1] if rolls == 1 else [1, 2, 0]
    else:
        raise NotImplementedError

    A_rolled_ref = np.transpose(A, perm)

    if direction == "left":
        A_rolled_tf = rollaxis_left(A_tf, rolls)
    elif direction == "right":
        A_rolled_tf = rollaxis_right(A_tf, rolls)
    else:
        raise NotImplementedError

    assert_allclose(A_rolled_ref, A_rolled_tf)


@pytest.mark.parametrize("rolls", [1, 2])
def test_rollaxis_idempotent(rolls: int) -> None:
    A = np.random.randn(10, 5, 3, 20, 1)
    A_tf = tf.convert_to_tensor(A)
    A_left_right = rollaxis_left(rollaxis_right(A_tf, 2), 2)
    A_right_left = rollaxis_right(rollaxis_left(A_tf, 2), 2)

    assert_allclose(A, A_left_right)
    assert_allclose(A, A_right_left)


@pytest.mark.parametrize("leading_dims", [tuple(), (1,), (5,)])
@pytest.mark.parametrize("n", [1, 5])
@pytest.mark.parametrize("d", [1, 5])
@pytest.mark.parametrize("num_samples", [None, 1, 5])
@pytest.mark.parametrize("full_cov", [True, False])
def test_sample_mvn_shapes(
    leading_dims: Tuple[int, ...], n: int, d: int, num_samples: Optional[int], full_cov: bool
) -> None:
    mean_shape = leading_dims + (n, d)
    means = tf.zeros(mean_shape, dtype=default_float())

    if full_cov:
        covariance_shape = leading_dims + (n, d, d)
        sqrt_cov = tf.random.normal(covariance_shape, dtype=default_float())
        covariances = tf.matmul(sqrt_cov, sqrt_cov, transpose_b=True)
    else:
        covariance_shape = leading_dims + (n, d)
        covariances = tf.random.normal(covariance_shape, dtype=default_float())

    samples = sample_mvn(means, covariances, full_cov, num_samples)

    if num_samples:
        expected_shape = leading_dims + (num_samples, n, d)
    else:
        expected_shape = leading_dims + (n, d)

    assert_equal(samples.shape, expected_shape)


# tests/gpflow/config/test_config.py

# Copyright 2018 the GPflow authors.
#
# 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.

import os
from typing import Any, Callable
from unittest import mock

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import TensorData
from gpflow.config import (
    default_float,
    default_int,
    default_jitter,
    default_positive_bijector,
    default_summary_fmt,
    set_default_float,
    set_default_int,
    set_default_jitter,
    set_default_positive_bijector,
    set_default_summary_fmt,
)
from gpflow.utilities import to_default_float, to_default_int

_env_values = [
    ("int", "int16", np.int16),
    ("int", "int64", np.int64),
    ("float", "float16", np.float16),
    ("float", "float32", np.float32),
    ("positive_bijector", "exp", "exp"),
    ("positive_bijector", "softplus", "softplus"),
    ("summary_fmt", "simple", "simple"),
    ("positive_minimum", "1e-3", 1e-3),
    ("jitter", "1e-2", 1e-2),
]


@pytest.mark.parametrize("attr_name, value, expected_value", _env_values)
def test_env_variables(attr_name: str, value: str, expected_value: Any) -> None:
    env_name = f"GPFLOW_{attr_name.upper()}"
    with mock.patch.dict("os.environ", {env_name: value}):
        assert os.environ[env_name] == value
        config = gpflow.config.Config()
        assert getattr(config, attr_name) == expected_value


@pytest.mark.parametrize("attr_name", dict.fromkeys(list(zip(*_env_values))[0]).keys())
def test_env_variables_failures(attr_name: str) -> None:
    if attr_name == "summary_fmt":
        pytest.skip("The `summary_fmt` validation cannot be performed.")
    env_name = f"GPFLOW_{attr_name.upper()}"
    with mock.patch.dict("os.environ", {env_name: "garbage"}):
        with pytest.raises(TypeError):
            gpflow.config.Config()


@pytest.mark.parametrize(
    "getter, setter, valid_type_1, valid_type_2",
    [
        (default_int, set_default_int, tf.int64, np.int32),
        (default_float, set_default_float, tf.float32, np.float64),
    ],
)
def test_dtype_setting(
    getter: Callable[[], type],
    setter: Callable[[type], None],
    valid_type_1: type,
    valid_type_2: type,
) -> None:
    if valid_type_1 == valid_type_2:
        raise ValueError("cannot test config setting/getting when both types are equal")
    setter(valid_type_1)
    assert getter() == valid_type_1
    setter(valid_type_2)
    assert getter() == valid_type_2


@pytest.mark.parametrize(
    "setter, invalid_type",
    [
        (set_default_int, str),
        (set_default_int, np.float64),
        (set_default_float, list),
        (set_default_float, tf.int32),
    ],
)
def test_dtype_errorcheck(setter: Callable[[type], None], invalid_type: Any) -> None:
    with pytest.raises(TypeError):
        setter(invalid_type)


def test_jitter_setting() -> None:
    set_default_jitter(1e-3)
    assert default_jitter() == 1e-3
    set_default_jitter(1e-6)
    assert default_jitter() == 1e-6


def test_jitter_errorcheck() -> None:
    with pytest.raises(TypeError):
        set_default_jitter("not a float")  # type: ignore[arg-type]
    with pytest.raises(ValueError):
        set_default_jitter(-1e-10)


@pytest.mark.parametrize(
    "value, error_msg",
    [
        ("Unknown", r"`unknown` not in set of valid bijectors: \['exp', 'softplus'\]"),
        (1.0, r"`1.0` not in set of valid bijectors: \['exp', 'softplus'\]"),
    ],
)
def test_positive_bijector_error(value: Any, error_msg: str) -> None:
    with pytest.raises(ValueError, match=error_msg):
        set_default_positive_bijector(value)


@pytest.mark.parametrize("value", ["exp", "SoftPlus"])
def test_positive_bijector_setting(value: str) -> None:
    set_default_positive_bijector(value)
    assert default_positive_bijector() == value.lower()


def test_default_summary_fmt_setting() -> None:
    set_default_summary_fmt("html")
    assert default_summary_fmt() == "html"
    set_default_summary_fmt(None)
    assert default_summary_fmt() is None


def test_default_summary_fmt_errorcheck() -> None:
    with pytest.raises(ValueError):
        set_default_summary_fmt("this_format_definitely_does_not_exist")


@pytest.mark.parametrize(
    "setter, getter, converter, dtype, value",
    [
        (set_default_int, default_int, to_default_int, np.int32, 3),
        (set_default_int, default_int, to_default_int, tf.int32, 3),
        (set_default_int, default_int, to_default_int, tf.int64, [3, 1, 4, 1, 5, 9]),
        (set_default_int, default_int, to_default_int, np.int64, [3, 1, 4, 1, 5, 9]),
        (set_default_float, default_float, to_default_float, np.float32, 3.14159),
        (
            set_default_float,
            default_float,
            to_default_float,
            tf.float32,
            [3.14159, 3.14159, 3.14159],
        ),
        (
            set_default_float,
            default_float,
            to_default_float,
            np.float64,
            [3.14159, 3.14159, 3.14159],
        ),
        (
            set_default_float,
            default_float,
            to_default_float,
            tf.float64,
            [3.14159, 3.14159, 3.14159],
        ),
    ],
)
def test_native_to_default_dtype(
    setter: Callable[[type], None],
    getter: Callable[[], type],
    converter: Callable[[TensorData], tf.Tensor],
    dtype: type,
    value: TensorData,
) -> None:
    with gpflow.config.as_context():
        setter(dtype)
        assert converter(value).dtype == dtype
        assert converter(value).dtype == getter()


# tests/gpflow/covariances/test_base_covariances.py

# Copyright 2017 Mark van der Wilk
#
# 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.

import numpy as np
import pytest
from numpy.testing import assert_allclose

import gpflow
from gpflow.config import default_jitter
from gpflow.covariances import Kuf, Kuu
from gpflow.inducing_variables import InducingPatches, InducingPoints, InducingVariables, Multiscale
from gpflow.kernels import Kernel


@pytest.mark.parametrize("N, D", [[17, 3], [10, 7]])
def test_inducing_points_inducing_variable_len(N: int, D: int) -> None:
    Z = np.random.randn(N, D)
    inducing_variable = InducingPoints(Z)
    assert inducing_variable.num_inducing == N


_kernel_setups = [
    gpflow.kernels.SquaredExponential(variance=0.46, lengthscales=np.random.uniform(0.5, 3.0, 5)),
    gpflow.kernels.Periodic(
        base_kernel=gpflow.kernels.SquaredExponential(variance=1.8), period=0.4
    ),
]


@pytest.mark.parametrize("N", [10, 101])
@pytest.mark.parametrize("kernel", _kernel_setups)
def test_inducing_equivalence(N: int, kernel: Kernel) -> None:
    # Inducing inducing must be the same as the kernel evaluations
    Z = np.random.randn(N, 5)
    inducing_variable = InducingPoints(Z)
    assert_allclose(Kuu(inducing_variable, kernel), kernel(Z))


@pytest.mark.parametrize("N, M, D", [[23, 13, 3], [10, 5, 7]])
def test_multi_scale_inducing_equivalence_inducing_points(N: int, M: int, D: int) -> None:
    # Multiscale must be equivalent to inducing points when variance is zero
    Xnew, Z = np.random.randn(N, D), np.random.randn(M, D)
    rbf = gpflow.kernels.SquaredExponential(1.3441, lengthscales=np.random.uniform(0.5, 3.0, D))
    inducing_variable_zero_lengthscales = Multiscale(Z, scales=np.zeros(Z.shape) + 1e-10)
    inducing_variable_inducing_point = InducingPoints(Z)

    multi_scale_Kuf = Kuf(inducing_variable_zero_lengthscales, rbf, Xnew)
    inducing_point_Kuf = Kuf(inducing_variable_inducing_point, rbf, Xnew)

    relative_error_Kuf = np.abs(multi_scale_Kuf - inducing_point_Kuf) / inducing_point_Kuf
    assert np.max(relative_error_Kuf) < 0.1e-2  # 0.1 %

    multi_scale_Kuu = Kuu(inducing_variable_zero_lengthscales, rbf)
    inducing_point_Kuu = Kuu(inducing_variable_inducing_point, rbf)

    relative_error_Kuu = np.abs(multi_scale_Kuu - inducing_point_Kuu) / inducing_point_Kuu
    assert np.max(relative_error_Kuu) < 0.1e-2  # 0.1 %


_inducing_variables_and_kernels = [
    [
        2,
        InducingPoints(np.random.randn(71, 2)),
        gpflow.kernels.SquaredExponential(
            variance=1.84, lengthscales=np.random.uniform(0.5, 3.0, 2)
        ),
    ],
    [
        2,
        InducingPoints(np.random.randn(71, 2)),
        gpflow.kernels.Matern12(variance=1.84, lengthscales=np.random.uniform(0.5, 3.0, 2)),
    ],
    [
        2,
        Multiscale(np.random.randn(71, 2), np.random.uniform(0.5, 3, size=(71, 2))),
        gpflow.kernels.SquaredExponential(
            variance=1.84, lengthscales=np.random.uniform(0.5, 3.0, 2)
        ),
    ],
    [
        9,
        InducingPatches(np.random.randn(71, 4)),
        gpflow.kernels.Convolutional(gpflow.kernels.SquaredExponential(), [3, 3], [2, 2]),
    ],
]


@pytest.mark.parametrize("input_dim, inducing_variable, kernel", _inducing_variables_and_kernels)
def test_inducing_variables_psd_schur(
    input_dim: int, inducing_variable: InducingVariables, kernel: Kernel
) -> None:
    # Conditional variance must be PSD.
    X = np.random.randn(5, input_dim)
    Kuf_values = Kuf(inducing_variable, kernel, X)
    Kuu_values = Kuu(inducing_variable, kernel, jitter=default_jitter())
    Kff_values = kernel(X)
    Qff_values = Kuf_values.numpy().T @ np.linalg.solve(Kuu_values, Kuf_values)
    assert np.all(np.linalg.eig(Kff_values - Qff_values)[0] > 0.0)


# tests/gpflow/covariances/test_multioutput.py

from typing import Any, Callable, Sequence

import numpy as np
import pytest
import tensorflow as tf

import gpflow
import gpflow.inducing_variables.multioutput as mf
import gpflow.kernels.multioutput as mk
from gpflow import covariances
from gpflow.inducing_variables import InducingVariables
from gpflow.kernels import Kernel

rng = np.random.RandomState(9911)

# ------------------------------------------
# Helpers
# ------------------------------------------


def make_kernel() -> Kernel:
    return gpflow.kernels.SquaredExponential()


def make_kernels(num: int) -> Sequence[Kernel]:
    return [make_kernel() for _ in range(num)]


def make_ip() -> InducingVariables:
    X = rng.permutation(Datum.X)
    return gpflow.inducing_variables.InducingPoints(X[: Datum.M, ...])


def make_ips(num: int) -> Sequence[InducingVariables]:
    return [make_ip() for _ in range(num)]


def type_name(x: Any) -> str:
    return x.__class__.__name__  # type: ignore[no-any-return]


# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


class Datum:
    D = 1
    L = 2
    P = 3
    M = 10
    N = 100
    W = rng.randn(P, L)
    X = rng.randn(N)[:, None]
    Xnew = rng.randn(N)[:, None]


multioutput_inducing_variable_list = [
    (
        mf.SharedIndependentInducingVariables(make_ip()),
        mf.SharedIndependentInducingVariables(make_ip()),
    ),
    (
        mf.SeparateIndependentInducingVariables(make_ips(Datum.P)),
        mf.SeparateIndependentInducingVariables(make_ips(Datum.L)),
    ),
]

multioutput_fallback_inducing_variable_list = [
    mf.FallbackSharedIndependentInducingVariables(make_ip()),
    mf.FallbackSeparateIndependentInducingVariables(make_ips(Datum.L)),
]

multioutput_kernel_list = [
    mk.SharedIndependent(make_kernel(), Datum.P),
    mk.SeparateIndependent(make_kernels(Datum.P)),
    mk.LinearCoregionalization(make_kernels(Datum.L), Datum.W),
]


@pytest.mark.parametrize(
    "inducing_variable_P,inducing_variable_L", multioutput_inducing_variable_list, ids=type_name
)
@pytest.mark.parametrize("kernel", multioutput_kernel_list, ids=type_name)
def test_kuu_shape(
    inducing_variable_P: InducingVariables, inducing_variable_L: InducingVariables, kernel: Kernel
) -> None:
    if isinstance(kernel, mk.LinearCoregionalization):
        inducing_variable = inducing_variable_L
    else:
        inducing_variable = inducing_variable_P

    Kuu = covariances.Kuu(inducing_variable, kernel, jitter=1e-9)
    t = tf.linalg.cholesky(Kuu)

    if isinstance(kernel, mk.SharedIndependent):
        if isinstance(inducing_variable, mf.SeparateIndependentInducingVariables):
            assert t.shape == (Datum.P, Datum.M, Datum.M)
        else:
            assert t.shape == (Datum.M, Datum.M)
    elif isinstance(kernel, mk.LinearCoregionalization):
        assert t.shape == (Datum.L, Datum.M, Datum.M)
    else:
        assert t.shape == (Datum.P, Datum.M, Datum.M)


@pytest.mark.parametrize(
    "inducing_variable_P,inducing_variable_L", multioutput_inducing_variable_list, ids=type_name
)
@pytest.mark.parametrize("kernel", multioutput_kernel_list, ids=type_name)
def test_kuf_shape(
    inducing_variable_P: InducingVariables, inducing_variable_L: InducingVariables, kernel: Kernel
) -> None:
    if isinstance(kernel, mk.LinearCoregionalization):
        inducing_variable = inducing_variable_L
    else:
        inducing_variable = inducing_variable_P

    Kuf = covariances.Kuf(inducing_variable, kernel, Datum.Xnew)

    if isinstance(kernel, mk.SharedIndependent):
        if isinstance(inducing_variable, mf.SeparateIndependentInducingVariables):
            assert Kuf.shape == (Datum.P, Datum.M, Datum.N)
        else:
            assert Kuf.shape == (Datum.M, Datum.N)
    elif isinstance(kernel, mk.LinearCoregionalization):
        assert Kuf.shape == (Datum.L, Datum.M, Datum.N)
    else:
        assert Kuf.shape == (Datum.P, Datum.M, Datum.N)


@pytest.mark.parametrize(
    "inducing_variable", multioutput_fallback_inducing_variable_list, ids=type_name
)
def test_kuf_fallback_shared_inducing_variables_shape(inducing_variable: InducingVariables) -> None:
    kernel = mk.LinearCoregionalization(make_kernels(Datum.L), Datum.W)
    Kuf = covariances.Kuf(inducing_variable, kernel, Datum.Xnew)

    assert Kuf.shape == (Datum.M, Datum.L, Datum.N, Datum.P)


@pytest.mark.parametrize("fun", [covariances.Kuu, covariances.Kuf], ids=type_name)
def test_mixed_shared_shape(fun: Callable[..., tf.Tensor]) -> None:
    inducing_variable = mf.SharedIndependentInducingVariables(make_ip())
    kernel = mk.LinearCoregionalization(make_kernels(Datum.L), Datum.W)
    if fun is covariances.Kuu:
        t = tf.linalg.cholesky(fun(inducing_variable, kernel, jitter=1e-9))
        assert t.shape == (Datum.L, Datum.M, Datum.M)
    else:
        t = fun(inducing_variable, kernel, Datum.Xnew)
        assert t.shape == (Datum.L, Datum.M, Datum.N)


# tests/gpflow/expectations/test_expectations.py

# Copyright 2018 the GPflow authors.
#
# 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 typing import Any, Callable, Iterable, Mapping, Sequence

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose
from tensorflow import convert_to_tensor as ctt

import gpflow
import gpflow.inducing_variables as iv
import gpflow.kernels as krn
from gpflow import mean_functions as mf
from gpflow.base import AnyNDArray
from gpflow.config import default_float
from gpflow.expectations import expectation, quadrature_expectation
from gpflow.probability_distributions import (
    DiagonalGaussian,
    Gaussian,
    MarkovGaussian,
    ProbabilityDistribution,
)

rng = np.random.RandomState(1)
RTOL = 1e-6

num_data = 5
num_ind = 4
D_in = 2
D_out = 2

Xmu = ctt(rng.randn(num_data, D_in))
Xmu_markov = ctt(rng.randn(num_data + 1, D_in))  # (N+1)xD
Xcov = rng.randn(num_data, D_in, D_in)
Xcov = ctt(Xcov @ np.transpose(Xcov, (0, 2, 1)))
Z = rng.randn(num_ind, D_in)


def markov_gauss() -> MarkovGaussian:
    cov_params: AnyNDArray = rng.randn(num_data + 1, D_in, 2 * D_in) / 2.0  # (N+1)xDx2D
    Xcov: AnyNDArray = cov_params @ np.transpose(cov_params, (0, 2, 1))  # (N+1)xDxD
    Xcross = cov_params[:-1] @ np.transpose(cov_params[1:], (0, 2, 1))  # NxDxD
    Xcross = np.concatenate((Xcross, np.zeros((1, D_in, D_in))), 0)  # (N+1)xDxD
    Xcov = np.stack([Xcov, Xcross])  # 2x(N+1)xDxD
    return MarkovGaussian(Xmu_markov, ctt(Xcov))


_means: Mapping[str, mf.MeanFunction] = {
    "lin": mf.Linear(A=rng.randn(D_in, D_out), b=rng.randn(D_out)),
    "identity": mf.Identity(input_dim=D_in),
    "const": mf.Constant(c=rng.randn(D_out)),
    "zero": mf.Zero(output_dim=D_out),
}

_distrs: Mapping[str, ProbabilityDistribution] = {
    "gauss": Gaussian(Xmu, Xcov),
    "dirac_gauss": Gaussian(Xmu, np.zeros((num_data, D_in, D_in))),
    "gauss_diag": DiagonalGaussian(Xmu, rng.rand(num_data, D_in)),
    "dirac_diag": DiagonalGaussian(Xmu, np.zeros((num_data, D_in))),
    "dirac_markov_gauss": MarkovGaussian(Xmu_markov, np.zeros((2, num_data + 1, D_in, D_in))),
    "markov_gauss": markov_gauss(),
}

_kerns: Mapping[str, krn.Kernel] = {
    "rbf": krn.SquaredExponential(variance=rng.rand(), lengthscales=rng.rand() + 1.0),
    "lin": krn.Linear(variance=rng.rand()),
    "matern": krn.Matern32(variance=rng.rand()),
    "rbf_act_dim_0": krn.SquaredExponential(
        variance=rng.rand(), lengthscales=rng.rand() + 1.0, active_dims=[0]
    ),
    "rbf_act_dim_1": krn.SquaredExponential(
        variance=rng.rand(), lengthscales=rng.rand() + 1.0, active_dims=[1]
    ),
    "lin_act_dim_0": krn.Linear(variance=rng.rand(), active_dims=[0]),
    "lin_act_dim_1": krn.Linear(variance=rng.rand(), active_dims=[1]),
    "rbf_lin_sum": krn.Sum(
        [
            krn.SquaredExponential(variance=rng.rand(), lengthscales=rng.rand() + 1.0),
            krn.Linear(variance=rng.rand()),
        ]
    ),
    "rbf_lin_sum2": krn.Sum(
        [
            krn.Linear(variance=rng.rand()),
            krn.SquaredExponential(variance=rng.rand(), lengthscales=rng.rand() + 1.0),
            krn.Linear(variance=rng.rand()),
            krn.SquaredExponential(variance=rng.rand(), lengthscales=rng.rand() + 1.0),
        ]
    ),
    "rbf_lin_prod": krn.Product(
        [
            krn.SquaredExponential(
                variance=rng.rand(), lengthscales=rng.rand() + 1.0, active_dims=[0]
            ),
            krn.Linear(variance=rng.rand(), active_dims=[1]),
        ]
    ),
}


def kerns(*args: str) -> Sequence[krn.Kernel]:
    return [_kerns[k] for k in args]


def distrs(*args: str) -> Sequence[ProbabilityDistribution]:
    return [_distrs[k] for k in args]


def means(*args: str) -> Sequence[mf.MeanFunction]:
    return [_means[k] for k in args]


@pytest.fixture
def inducing_variable() -> iv.InducingVariables:
    return iv.InducingPoints(Z)


def _check(params: Iterable[Any]) -> None:
    analytic = expectation(*params)
    quad = quadrature_expectation(*params)
    assert_allclose(analytic, quad, rtol=RTOL)


# =================================== TESTS ===================================

distr_args1 = distrs("gauss")
mean_args = means("lin", "identity", "const", "zero")
kern_args1 = kerns("lin", "rbf", "rbf_lin_sum", "rbf_lin_prod")
kern_args2 = kerns("lin", "rbf", "rbf_lin_sum")


@pytest.mark.parametrize("distribution", distr_args1)
@pytest.mark.parametrize("mean1", mean_args)
@pytest.mark.parametrize("mean2", mean_args)
@pytest.mark.parametrize("arg_filter", [lambda p, m1, m2: (p, m1), lambda p, m1, m2: (p, m1, m2)])
def test_mean_function_only_expectations(
    distribution: ProbabilityDistribution,
    mean1: mf.MeanFunction,
    mean2: mf.MeanFunction,
    arg_filter: Callable[
        [ProbabilityDistribution, mf.MeanFunction, mf.MeanFunction], Iterable[Any]
    ],
) -> None:
    params = arg_filter(distribution, mean1, mean2)
    _check(params)


@pytest.mark.parametrize("distribution", distrs("gauss", "gauss_diag"))
@pytest.mark.parametrize("kernel", kern_args1)
@pytest.mark.parametrize(
    "arg_filter",
    [
        lambda p, k, f: (p, k),
        lambda p, k, f: (p, (k, f)),
        lambda p, k, f: (p, (k, f), (k, f)),
    ],
)
def test_kernel_only_expectations(
    distribution: ProbabilityDistribution,
    kernel: krn.Kernel,
    inducing_variable: iv.InducingVariables,
    arg_filter: Callable[[ProbabilityDistribution, krn.Kernel, mf.MeanFunction], Iterable[Any]],
) -> None:
    params = arg_filter(distribution, kernel, inducing_variable)
    _check(params)


@pytest.mark.parametrize("distribution", distr_args1)
@pytest.mark.parametrize("kernel", kerns("rbf", "lin", "matern", "rbf_lin_sum"))
@pytest.mark.parametrize("mean", mean_args)
@pytest.mark.parametrize(
    "arg_filter", [lambda p, k, f, m: (p, (k, f), m), lambda p, k, f, m: (p, m, (k, f))]
)
def test_kernel_mean_function_expectations(
    distribution: ProbabilityDistribution,
    kernel: krn.Kernel,
    inducing_variable: iv.InducingVariables,
    mean: mf.MeanFunction,
    arg_filter: Callable[
        [ProbabilityDistribution, krn.Kernel, iv.InducingVariables, mf.MeanFunction], Iterable[Any]
    ],
) -> None:
    params = arg_filter(distribution, kernel, inducing_variable, mean)
    _check(params)


@pytest.mark.parametrize("kernel", kern_args1)
def test_eKdiag_no_uncertainty(kernel: krn.Kernel) -> None:
    eKdiag = expectation(_distrs["dirac_diag"], kernel)
    Kdiag = kernel(Xmu, full_cov=False)
    assert_allclose(eKdiag, Kdiag, rtol=RTOL)


@pytest.mark.parametrize("kernel", kern_args1)
def test_eKxz_no_uncertainty(kernel: krn.Kernel, inducing_variable: iv.InducingVariables) -> None:
    eKxz = expectation(_distrs["dirac_diag"], (kernel, inducing_variable))
    Kxz = kernel(Xmu, Z)
    assert_allclose(eKxz, Kxz, rtol=RTOL)


@pytest.mark.parametrize("kernel", kern_args2)
@pytest.mark.parametrize("mean", mean_args)
def test_eMxKxz_no_uncertainty(
    kernel: krn.Kernel, inducing_variable: iv.InducingVariables, mean: mf.MeanFunction
) -> None:
    exKxz = expectation(_distrs["dirac_diag"], mean, (kernel, inducing_variable))
    Kxz = kernel(Xmu, Z)
    xKxz = expectation(_distrs["dirac_gauss"], mean)[:, :, None] * Kxz[:, None, :]
    assert_allclose(exKxz, xKxz, rtol=RTOL)


@pytest.mark.parametrize("kernel", kern_args1)
def test_eKzxKxz_no_uncertainty(
    kernel: krn.Kernel, inducing_variable: iv.InducingVariables
) -> None:
    eKzxKxz = expectation(
        _distrs["dirac_diag"], (kernel, inducing_variable), (kernel, inducing_variable)
    )
    Kxz = kernel(Xmu, Z)
    KzxKxz = Kxz[:, :, None] * Kxz[:, None, :]
    assert_allclose(eKzxKxz, KzxKxz, rtol=RTOL)


def test_RBF_eKzxKxz_gradient_notNaN() -> None:
    """
    Ensure that <K_{Z, x} K_{x, Z}>_p(x) is not NaN and correct, when
    K_{Z, Z} is zero with finite precision. See pull request #595.
    """
    kernel = krn.SquaredExponential(1, lengthscales=0.1)
    kernel.variance.assign(2.0)

    p = gpflow.probability_distributions.Gaussian(
        tf.constant([[10]], dtype=default_float()),
        tf.constant([[[0.1]]], dtype=default_float()),
    )
    z = iv.InducingPoints([[-10.0], [10.0]])

    with tf.GradientTape() as tape:
        ekz = expectation(p, (kernel, z), (kernel, z))
    grad = tape.gradient(ekz, kernel.lengthscales.unconstrained_variable)
    assert grad is not None and not np.isnan(grad)


@pytest.mark.parametrize("distribution", distrs("gauss_diag"))
@pytest.mark.parametrize("kern1", kerns("rbf_act_dim_0", "lin_act_dim_0"))
@pytest.mark.parametrize("kern2", kerns("rbf_act_dim_1", "lin_act_dim_1"))
def test_eKzxKxz_separate_dims_simplification(
    distribution: ProbabilityDistribution,
    kern1: krn.Kernel,
    kern2: krn.Kernel,
    inducing_variable: iv.InducingVariables,
) -> None:
    _check((distribution, (kern1, inducing_variable), (kern2, inducing_variable)))


@pytest.mark.parametrize("distribution", distr_args1)
@pytest.mark.parametrize("kern1", kerns("rbf_lin_sum"))
@pytest.mark.parametrize("kern2", kerns("rbf_lin_sum2"))
def test_eKzxKxz_different_sum_kernels(
    distribution: ProbabilityDistribution,
    kern1: krn.Kernel,
    kern2: krn.Kernel,
    inducing_variable: iv.InducingVariables,
) -> None:
    _check((distribution, (kern1, inducing_variable), (kern2, inducing_variable)))


@pytest.mark.parametrize("distribution", distr_args1)
@pytest.mark.parametrize("kern1", kerns("rbf_lin_sum2"))
@pytest.mark.parametrize("kern2", kerns("rbf_lin_sum2"))
def test_eKzxKxz_same_vs_different_sum_kernels(
    distribution: ProbabilityDistribution,
    kern1: krn.Kernel,
    kern2: krn.Kernel,
    inducing_variable: iv.InducingVariables,
) -> None:
    # check the result is the same if we pass different objects with the same value
    same = expectation(*(distribution, (kern1, inducing_variable), (kern1, inducing_variable)))
    different = expectation(*(distribution, (kern1, inducing_variable), (kern2, inducing_variable)))
    assert_allclose(same, different, rtol=RTOL)


@pytest.mark.parametrize("distribution", distrs("markov_gauss"))
@pytest.mark.parametrize("kernel", kern_args2)
@pytest.mark.parametrize("mean", means("identity"))
def test_exKxz_markov(
    distribution: ProbabilityDistribution,
    kernel: krn.Kernel,
    mean: mf.MeanFunction,
    inducing_variable: iv.InducingVariables,
) -> None:
    _check((distribution, (kernel, inducing_variable), mean))


@pytest.mark.parametrize("distribution", distrs("dirac_markov_gauss"))
@pytest.mark.parametrize("kernel", kern_args2)
@pytest.mark.parametrize("mean", means("identity"))
def test_exKxz_markov_no_uncertainty(
    distribution: ProbabilityDistribution,
    kernel: krn.Kernel,
    mean: mf.MeanFunction,
    inducing_variable: iv.InducingVariables,
) -> None:
    exKxz = expectation(distribution, (kernel, inducing_variable), mean)
    Kzx = kernel(Xmu_markov[:-1, :], Z)  # NxM
    xKxz = Kzx[..., None] * Xmu_markov[1:, None, :]  # NxMxD
    assert_allclose(exKxz, xKxz, rtol=RTOL)


@pytest.mark.parametrize("kernel", kerns("rbf"))
@pytest.mark.parametrize("distribution", distrs("gauss", "gauss_diag", "markov_gauss"))
def test_cov_shape_inference(
    distribution: ProbabilityDistribution,
    kernel: krn.Kernel,
    inducing_variable: iv.InducingVariables,
) -> None:
    assert isinstance(distribution, (Gaussian, DiagonalGaussian, MarkovGaussian))
    gauss_tuple = (distribution.mu, distribution.cov)
    _check((gauss_tuple, (kernel, inducing_variable)))
    if isinstance(distribution, MarkovGaussian):
        _check((gauss_tuple, None, (kernel, inducing_variable)))


# tests/gpflow/experimental/test_utils.py

# Copyright 2022 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 _pytest.recwarn import WarningsRecorder

from gpflow.experimental.utils import experimental


def test_experimental(recwarn: WarningsRecorder) -> None:
    @experimental
    def f() -> None:
        pass

    assert len(recwarn) == 0  # pylint: disable=len-as-condition
    f()
    assert len(recwarn) == 1
    f()
    assert len(recwarn) == 1


# tests/gpflow/kernels/reference.py

import numpy as np

from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import check_shapes


@check_shapes(
    "X: [N, D]",
    "X: [N, D_tuple...]",
    "lengthscales: [broadcast D_tuple...]",
    "signal_variance: [broadcast D_tuple...]",
)
def ref_rbf_kernel(
    X: AnyNDArray, lengthscales: AnyNDArray, signal_variance: AnyNDArray
) -> AnyNDArray:
    N, _ = X.shape
    kernel = np.zeros((N, N))
    for row_index in range(N):
        for column_index in range(N):
            vecA = X[row_index, :]
            vecB = X[column_index, :]
            delta = vecA - vecB
            distance_squared = np.dot(delta.T, delta)
            kernel[row_index, column_index] = signal_variance * np.exp(
                -0.5 * distance_squared / lengthscales ** 2
            )
    return kernel


@check_shapes(
    "X: [N, D]",
    "X: [N, D_tuple...]",
    "weight_variances: [broadcast D_tuple...]",
    "bias_variance: [broadcast D_tuple...]",
    "signal_variance: [broadcast D_tuple...]",
)
def ref_arccosine_kernel(
    X: AnyNDArray,
    order: int,
    weight_variances: AnyNDArray,
    bias_variance: AnyNDArray,
    signal_variance: AnyNDArray,
) -> AnyNDArray:
    num_points = X.shape[0]
    kernel = np.empty((num_points, num_points))
    for row in range(num_points):
        for col in range(num_points):
            x = X[row]
            y = X[col]

            numerator = (weight_variances * x).dot(y) + bias_variance

            x_denominator = np.sqrt((weight_variances * x).dot(x) + bias_variance)
            y_denominator = np.sqrt((weight_variances * y).dot(y) + bias_variance)
            denominator = x_denominator * y_denominator

            theta = np.arccos(np.clip(numerator / denominator, -1.0, 1.0))
            if order == 0:
                J = np.pi - theta
            elif order == 1:
                J = np.sin(theta) + (np.pi - theta) * np.cos(theta)
            elif order == 2:
                J = 3.0 * np.sin(theta) * np.cos(theta)
                J += (np.pi - theta) * (1.0 + 2.0 * np.cos(theta) ** 2)

            kernel[row, col] = (
                signal_variance
                * (1.0 / np.pi)
                * J
                * x_denominator ** order
                * y_denominator ** order
            )
    return kernel


@check_shapes(
    "X: [N, D]",
    "X: [N, D_tuple...]",
    "lengthscales: [broadcast D_tuple...]",
    "signal_variance: [broadcast D_tuple...]",
    "period: [broadcast D_tuple...]",
)
def ref_periodic_kernel(
    X: AnyNDArray,
    base_name: str,
    lengthscales: AnyNDArray,
    signal_variance: AnyNDArray,
    period: AnyNDArray,
) -> AnyNDArray:
    """
    Calculates K(X) for the periodic kernel based on various base kernels.
    """
    sine_arg = np.pi * (X[:, None, :] - X[None, :, :]) / period
    sine_base = np.sin(sine_arg) / lengthscales
    exp_dist: AnyNDArray
    if base_name in {"RBF", "SquaredExponential"}:
        dist = 0.5 * np.sum(np.square(sine_base), axis=-1)
        exp_dist = np.exp(-dist)
    elif base_name == "Matern12":
        dist = np.sum(np.abs(sine_base), axis=-1)
        exp_dist = np.exp(-dist)
    elif base_name == "Matern32":
        dist = np.sqrt(3) * np.sum(np.abs(sine_base), axis=-1)
        exp_dist = (1 + dist) * np.exp(-dist)
    elif base_name == "Matern52":
        dist = np.sqrt(5) * np.sum(np.abs(sine_base), axis=-1)
        exp_dist = (1 + dist + dist ** 2 / 3) * np.exp(-dist)
    return signal_variance * exp_dist


# tests/gpflow/kernels/test_broadcasting.py

# Copyright 2020 the GPflow authors.
#
# 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 typing import Any, Callable, List, Sequence, Tuple, cast

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
import gpflow.ci_utils
from gpflow import kernels
from gpflow.base import AnyNDArray, TensorType
from gpflow.experimental.check_shapes import check_shape as cs
from gpflow.experimental.check_shapes import check_shapes


def create_kernels() -> Sequence[kernels.Kernel]:
    result: List[kernels.Kernel] = [
        # Static kernels:
        kernels.White(),
        kernels.Constant(),
        # Stationary kernels:
        kernels.SquaredExponential(),
        kernels.RationalQuadratic(),
        kernels.Exponential(),
        kernels.Matern12(),
        kernels.Matern32(),
        kernels.Matern52(),
        # sum and product kernels:
        kernels.White() + kernels.Matern12(),
        kernels.White() * kernels.Matern12(),
        # slicing:
        kernels.Matern32(active_dims=slice(None, None, 2)),
        kernels.Matern32(active_dims=[1, 2]),
        # other kernels:
        kernels.Cosine(),
        kernels.Linear(),
        kernels.Polynomial(),
        kernels.Periodic(kernels.Matern32()),
        kernels.ChangePoints(
            [kernels.Matern32(), kernels.Matern32()],
            [0.5],
        ),
        kernels.ArcCosine(),
        kernels.Coregion(output_dim=5, rank=2),
        kernels.Convolutional(kernels.Matern32(), [4, 4], [2, 2]),
        # multioutput:
        kernels.SharedIndependent(kernels.Matern32(), output_dim=5),
        kernels.SeparateIndependent([kernels.Matern32() for _ in range(5)]),
        kernels.LinearCoregionalization([kernels.Matern32() for _ in range(3)], np.ones((5, 3))),
    ]
    return result


def make_id(value: Any) -> str:
    if isinstance(value, tuple):
        return f"[{','.join(repr(x) for x in value)}]"
    assert isinstance(value, kernels.Kernel)
    return value.__class__.__name__


def test_no_kernels_missed() -> None:
    tested_kernel_classes = {
        parent
        for kernel in create_kernels()
        for parent in kernel.__class__.__mro__
        if (parent != kernels.Kernel) and issubclass(parent, kernels.Kernel)
    }
    all_kernel_classes = set(gpflow.ci_utils.subclasses(kernels.Kernel))
    assert tested_kernel_classes == all_kernel_classes


@pytest.mark.parametrize("kernel", create_kernels(), ids=make_id)
@pytest.mark.parametrize(
    "batch_shape",
    [
        (3,),
        (2, 3),
        (1, 2, 3),
    ],
    ids=make_id,
)
@pytest.mark.parametrize(
    "batch2_shape",
    [
        (4,),
        (2, 4),
        (1, 2, 4),
    ],
    ids=make_id,
)
@check_shapes()
def test_broadcasting(
    kernel: kernels.Kernel,
    batch_shape: Tuple[int, ...],
    batch2_shape: Tuple[int, ...],
    rng: np.random.Generator,
) -> None:
    if isinstance(kernel, kernels.Coregion):
        D = 1
        X: AnyNDArray = rng.choice(kernel.rank, batch_shape + (D,))
        X2: AnyNDArray = rng.choice(kernel.rank, batch2_shape + (D,))
    else:
        if isinstance(kernel, kernels.ChangePoints):
            D = 1
        elif isinstance(kernel, kernels.Convolutional):
            D = int(np.prod(kernel.image_shape))
        else:
            D = 4

        X = rng.random(batch_shape + (D,))
        X2 = rng.random(batch2_shape + (D,))

    rank = len(batch_shape) - 1
    rank2 = len(batch2_shape) - 1

    if isinstance(kernel, kernels.MultioutputKernel):
        mo_kernel = cast(kernels.MultioutputKernel, kernel)

        loop = cs(
            unroll_batches(
                lambda x: unroll_batches(
                    lambda x2: mo_kernel(x, x2, full_cov=True, full_output_cov=True), X2, 2
                ),
                X,
                2,
            ),
            "[batch..., batch2..., N, P, N2, P]",
        )
        loop = cs(
            tf.transpose(
                loop,
                tf.concat(
                    [
                        np.arange(rank),
                        [rank + rank2, rank + rank2 + 1],
                        np.arange(rank2) + rank,
                        [rank + rank2 + 2, rank + rank2 + 3],
                    ],
                    0,
                ),
            ),
            "[batch..., N, P, batch2..., N2, P]",
        )
        native = cs(
            mo_kernel(X, X2, full_cov=True, full_output_cov=True),
            "[batch..., N, P, batch2..., N2, P]",
        )
        assert_allclose(loop.numpy(), native.numpy())

        loop = cs(
            unroll_batches(
                lambda x: unroll_batches(
                    lambda x2: mo_kernel(x, x2, full_cov=True, full_output_cov=False), X2, 2
                ),
                X,
                2,
            ),
            "[batch..., batch2..., P, N, N2]",
        )
        loop = cs(
            tf.transpose(
                loop,
                tf.concat(
                    [
                        [rank + rank2],
                        np.arange(rank),
                        [rank + rank2 + 1],
                        np.arange(rank2) + rank,
                        [rank + rank2 + 2],
                    ],
                    0,
                ),
            ),
            "[P, batch..., N, batch2..., N2]",
        )
        native = cs(
            mo_kernel(X, X2, full_cov=True, full_output_cov=False),
            "[P, batch..., N, batch2..., N2]",
        )
        assert_allclose(loop.numpy(), native.numpy())

        loop = cs(
            unroll_batches(lambda x: mo_kernel(x, full_cov=True, full_output_cov=True), X, 2),
            "[batch..., N, P, N, P]",
        )
        native = cs(mo_kernel(X, full_cov=True, full_output_cov=True), "[batch..., N, P, N, P]")
        assert_allclose(loop.numpy(), native.numpy())

        loop = cs(
            unroll_batches(lambda x: mo_kernel(x, full_cov=True, full_output_cov=False), X, 2),
            "[batch..., P, N, N]",
        )
        loop = cs(
            tf.transpose(
                loop,
                tf.concat(
                    [
                        [rank],
                        np.arange(rank),
                        [rank + 1, rank + 2],
                    ],
                    0,
                ),
            ),
            "[P, batch..., N, N]",
        )
        native = cs(mo_kernel(X, full_cov=True, full_output_cov=False), "[P, batch..., N, N]")
        assert_allclose(loop.numpy(), native.numpy())

        loop = cs(
            unroll_batches(lambda x: mo_kernel(x, full_cov=False, full_output_cov=True), X, 2),
            "[batch..., N, P, P]",
        )
        native = cs(mo_kernel(X, full_cov=False, full_output_cov=True), "[batch..., N, P, P]")
        assert_allclose(loop.numpy(), native.numpy())

        loop = cs(
            unroll_batches(lambda x: mo_kernel(x, full_cov=False, full_output_cov=False), X, 2),
            "[batch..., N, P]",
        )
        native = cs(mo_kernel(X, full_cov=False, full_output_cov=False), "[batch..., N, P]")
        assert_allclose(loop.numpy(), native.numpy())

    else:  # Single-output kernel:
        loop = cs(
            unroll_batches(
                lambda x: unroll_batches(lambda x2: kernel(x, x2, full_cov=True), X2, 2), X, 2
            ),
            "[batch..., batch2..., N, N2]",
        )
        loop = cs(
            tf.transpose(
                loop,
                tf.concat(
                    [np.arange(rank), [rank + rank2], np.arange(rank2) + rank, [rank + rank2 + 1]],
                    0,
                ),
            ),
            "[batch..., N, batch2..., N2]",
        )
        native = cs(kernel(X, X2, full_cov=True), "[batch..., N, batch2..., N2]")
        assert_allclose(loop.numpy(), native.numpy())

        loop = cs(unroll_batches(lambda x: kernel(x, full_cov=True), X, 2), "[batch..., N, N]")
        native = cs(kernel(X, full_cov=True), "[batch..., N, N]")
        assert_allclose(loop.numpy(), native.numpy())

        loop = cs(unroll_batches(lambda x: kernel(x, full_cov=False), X, 2), "[batch..., N]")
        native = cs(kernel(X, full_cov=False), "[batch..., N]")
        assert_allclose(loop.numpy(), native.numpy())


@check_shapes(
    "x: [batch..., value_shape...]",
    "return: [batch..., result_shape...]",
)
def unroll_batches(
    f: Callable[[TensorType], TensorType], x: TensorType, value_rank: int
) -> tf.Tensor:
    if len(x.shape) == value_rank:
        return f(x)

    return tf.stack([unroll_batches(f, row, value_rank) for row in x])


# tests/gpflow/kernels/test_changepoints.py

import numpy as np

import gpflow


def test_changepoint_with_X1_X2() -> None:
    N = 100
    X = np.linspace(0, 100, N).reshape(N, 1)
    base_k1 = gpflow.kernels.Matern32(lengthscales=0.2)
    base_k2 = gpflow.kernels.Matern32(lengthscales=2.0)
    k = gpflow.kernels.ChangePoints([base_k1, base_k2], [0.0], steepness=5.0)
    K = k(X)
    assert K.shape == [N, N]

    N2 = 25
    X2 = np.linspace(0, 50, N2).reshape(N2, 1)
    K = k(X, X2)
    assert K.shape == [N, N2]


# tests/gpflow/kernels/test_coregion.py

# Copyright 2017 the GPflow authors.
#
# 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 typing import Sequence, Tuple

import numpy as np
from numpy.testing import assert_allclose

import gpflow
from gpflow import set_trainable
from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.mean_functions import Constant
from gpflow.models import VGP

rng = np.random.RandomState(0)


class Datum:
    cs = ShapeChecker().check_shape

    N1, N2 = 6, 16
    X: Sequence[AnyNDArray] = [
        cs(rng.rand(N1, 2) * 1, "[N1, 2]"),
        cs(rng.rand(N2, 2) * 1, "[N2, 2]"),
    ]
    Y = [
        np.sin(x[:, :1]) + 0.9 * np.cos(x[:, 1:2] * 1.6) + rng.randn(x.shape[0], 1) * 0.8 for x in X
    ]
    label: Sequence[AnyNDArray] = [np.zeros((N1, 1)), np.ones((N2, 1))]
    X_augmented0: AnyNDArray = cs(np.hstack([X[0], label[0]]), "[N1, 3]")
    X_augmented1: AnyNDArray = cs(np.hstack([X[1], label[1]]), "[N2, 3]")
    X_augmented: AnyNDArray = cs(np.vstack([X_augmented0, X_augmented1]), "[N3, 3]")

    Y_augmented0: AnyNDArray = cs(np.hstack([Y[0], label[0]]), "[N1, 2]")
    Y_augmented1: AnyNDArray = cs(np.hstack([Y[1], label[1]]), "[N2, 2]")
    Y_augmented: AnyNDArray = cs(np.vstack([Y_augmented0, Y_augmented1]), "[N3, 2]")

    # For predict tests
    N = 10
    Xtest: AnyNDArray = rng.rand(N, 2) * N
    Xtest_augmented0: AnyNDArray = cs(np.hstack([Xtest, np.zeros((N, 1))]), "[N, 3]")
    Xtest_augmented1: AnyNDArray = cs(np.hstack([Xtest, np.ones((N, 1))]), "[N, 3]")
    Ytest = cs(np.sin(Xtest[:, :1]) + 0.9 * np.cos(Xtest[:, 1:2] * 1.6), "[N, 1]")
    Ytest_augmented0: AnyNDArray = cs(np.hstack([Ytest, np.zeros((N, 1))]), "[N, 2]")
    Ytest_augmented1: AnyNDArray = cs(np.hstack([Ytest, np.ones((N, 1))]), "[N, 2]")


def _prepare_models() -> Tuple[VGP, VGP, VGP]:
    """
    Prepare models to make sure the coregionalized model with diagonal coregion kernel and
    with fixed lengthscales is equivalent with normal GP regression.
    """
    # 1. Two independent VGPs for two sets of data
    k0 = gpflow.kernels.SquaredExponential()
    set_trainable(k0.lengthscales, False)
    k1 = gpflow.kernels.SquaredExponential()
    set_trainable(k1.lengthscales, False)
    vgp0 = VGP(
        (Datum.X[0], Datum.Y[0]),
        kernel=k0,
        mean_function=Constant(),
        likelihood=gpflow.likelihoods.Gaussian(),
        num_latent_gps=1,
    )
    vgp1 = VGP(
        (Datum.X[1], Datum.Y[1]),
        kernel=k1,
        mean_function=Constant(),
        likelihood=gpflow.likelihoods.Gaussian(),
        num_latent_gps=1,
    )
    # 2. Coregionalized VGP
    kc = gpflow.kernels.SquaredExponential(active_dims=[0, 1])
    set_trainable(kc.lengthscales, False)
    set_trainable(kc.variance, False)  # variance is handled by the Coregion kernel
    coreg = gpflow.kernels.Coregion(output_dim=2, rank=1, active_dims=[2])
    coreg.W.assign(np.zeros((2, 1)))  # zero correlation between outputs
    set_trainable(coreg.W, False)
    lik = gpflow.likelihoods.SwitchedLikelihood(
        [gpflow.likelihoods.Gaussian(), gpflow.likelihoods.Gaussian()]
    )
    mean_c = gpflow.mean_functions.SwitchedMeanFunction(
        [gpflow.mean_functions.Constant(), gpflow.mean_functions.Constant()]
    )
    cvgp = VGP(
        (Datum.X_augmented, Datum.Y_augmented),
        kernel=kc * coreg,
        mean_function=mean_c,
        likelihood=lik,
        num_latent_gps=1,
    )

    # Train them for a small number of iterations

    opt = gpflow.optimizers.Scipy()
    opt.minimize(
        vgp0.training_loss,
        variables=vgp0.trainable_variables,
        options=dict(maxiter=1000),
        method="BFGS",
    )
    opt.minimize(
        vgp1.training_loss,
        variables=vgp1.trainable_variables,
        options=dict(maxiter=1000),
        method="BFGS",
    )
    opt.minimize(
        cvgp.training_loss,
        variables=cvgp.trainable_variables,
        options=dict(maxiter=1000),
        method="BFGS",
    )

    return vgp0, vgp1, cvgp


# ------------------------------------------
# Tests
# ------------------------------------------


def test_likelihood_variance() -> None:
    vgp0, vgp1, cvgp = _prepare_models()
    assert_allclose(
        vgp0.likelihood.variance.numpy(),
        cvgp.likelihood.likelihoods[0].variance.numpy(),
        atol=1e-2,
    )
    assert_allclose(
        vgp1.likelihood.variance.numpy(),
        cvgp.likelihood.likelihoods[1].variance.numpy(),
        atol=1e-2,
    )


def test_kernel_variance() -> None:
    vgp0, vgp1, cvgp = _prepare_models()
    assert_allclose(
        vgp0.kernel.variance.numpy(),
        cvgp.kernel.kernels[1].kappa.numpy()[0],
        atol=1.0e-4,
    )
    assert_allclose(
        vgp1.kernel.variance.numpy(),
        cvgp.kernel.kernels[1].kappa.numpy()[1],
        atol=1.0e-4,
    )


def test_mean_values() -> None:
    vgp0, vgp1, cvgp = _prepare_models()
    assert_allclose(
        vgp0.mean_function.c.numpy(),
        cvgp.mean_function.meanfunctions[0].c.numpy(),
        atol=1.0e-4,
    )
    assert_allclose(
        vgp1.mean_function.c.numpy(),
        cvgp.mean_function.meanfunctions[1].c.numpy(),
        atol=1.0e-4,
    )


def test_predict_f() -> None:
    vgp0, vgp1, cvgp = _prepare_models()

    pred_f0 = vgp0.predict_f(Datum.Xtest)
    pred_fc0 = cvgp.predict_f(Datum.Xtest_augmented0)
    assert_allclose(pred_f0, pred_fc0, atol=1.0e-4)
    pred_f1 = vgp1.predict_f(Datum.Xtest)
    pred_fc1 = cvgp.predict_f(Datum.Xtest_augmented1)
    assert_allclose(pred_f1, pred_fc1, atol=1.0e-4)

    # check predict_f_full_cov
    vgp0.predict_f(Datum.Xtest, full_cov=True)
    cvgp.predict_f(Datum.Xtest_augmented0, full_cov=True)
    vgp1.predict_f(Datum.Xtest, full_cov=True)
    cvgp.predict_f(Datum.Xtest_augmented1, full_cov=True)


def test_predict_y() -> None:
    vgp0, vgp1, cvgp = _prepare_models()
    mu1, var1 = vgp0.predict_y(Datum.Xtest)
    c_mu1, c_var1 = cvgp.predict_y(np.hstack([Datum.Xtest, np.zeros((Datum.Xtest.shape[0], 1))]))

    # predict_y returns results for all the likelihoods in multi_likelihood
    assert_allclose(mu1, c_mu1[:, :1], atol=1.0e-4)
    assert_allclose(var1, c_var1[:, :1], atol=1.0e-4)

    mu2, var2 = vgp1.predict_y(Datum.Xtest)
    c_mu2, c_var2 = cvgp.predict_y(np.hstack([Datum.Xtest, np.ones((Datum.Xtest.shape[0], 1))]))

    # predict_y returns results for all the likelihoods in multi_likelihood
    assert_allclose(mu2, c_mu2[:, 1:2], atol=1.0e-4)
    assert_allclose(var2, c_var2[:, 1:2], atol=1.0e-4)


def test_predict_log_density() -> None:
    vgp0, vgp1, cvgp = _prepare_models()

    pred_ydensity0 = vgp0.predict_log_density((Datum.Xtest, Datum.Ytest))
    pred_ydensity_c0 = cvgp.predict_log_density((Datum.Xtest_augmented0, Datum.Ytest_augmented0))
    assert_allclose(pred_ydensity0, pred_ydensity_c0, atol=1e-2)
    pred_ydensity1 = vgp1.predict_log_density((Datum.Xtest, Datum.Ytest))
    pred_ydensity_c1 = cvgp.predict_log_density((Datum.Xtest_augmented1, Datum.Ytest_augmented1))
    assert_allclose(pred_ydensity1, pred_ydensity_c1, atol=1e-2)


def test_predict_f_samples() -> None:
    vgp0, vgp1, cvgp = _prepare_models()
    # just check predict_f_samples(self) works
    cvgp.predict_f_samples(Datum.X_augmented0, 1)
    cvgp.predict_f_samples(Datum.X_augmented1, 1)


# tests/gpflow/kernels/test_kernels.py

# Copyright 2018 the GPflow authors.
#
# 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 typing import Any, Optional, Sequence, Tuple, Type, Union, cast

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow.ci_utils
from gpflow.base import AnyNDArray, TensorType
from gpflow.config import default_float
from gpflow.experimental.check_shapes import check_shapes
from gpflow.kernels import (
    RBF,
    AnisotropicStationary,
    ArcCosine,
    ChangePoints,
    Constant,
    Convolutional,
    Coregion,
    Cosine,
    IsotropicStationary,
    Kernel,
    Linear,
    LinearCoregionalization,
    Matern12,
    Matern32,
    Matern52,
    MultioutputKernel,
    Periodic,
    Polynomial,
    RationalQuadratic,
    SeparateIndependent,
    SharedIndependent,
    SquaredExponential,
    Stationary,
    White,
)
from tests.gpflow.kernels.reference import ref_arccosine_kernel, ref_periodic_kernel, ref_rbf_kernel

rng = np.random.RandomState(1)


@check_shapes(
    "X: [N, D]",
    "locations: [L]",
    "steepness: [broadcast L]",
)
def _ref_changepoints(
    X: AnyNDArray,
    kernels: Sequence[Kernel],
    locations: Sequence[float],
    steepness: Union[float, Sequence[float]],
) -> AnyNDArray:
    """
    Calculates K(X) for each kernel in `kernels`, then multiply by sigmoid functions
    in order to smoothly transition betwen them. The sigmoid transitions are defined
    by a location and a steepness parameter.
    """
    locations = sorted(locations)
    steepness = steepness if isinstance(steepness, Sequence) else [steepness] * len(locations)
    np_locations: AnyNDArray = np.array(locations).reshape((1, 1, -1))
    np_steepness: AnyNDArray = np.array(steepness).reshape((1, 1, -1))

    sig_X = 1.0 / (1.0 + np.exp(-np_steepness * (X[:, :, None] - np_locations)))

    starters = sig_X * np.transpose(sig_X, axes=(1, 0, 2))
    stoppers = (1 - sig_X) * np.transpose((1 - sig_X), axes=(1, 0, 2))

    ones = np.ones((X.shape[0], X.shape[0], 1))
    starters = np.concatenate([ones, starters], axis=2)
    stoppers = np.concatenate([stoppers, ones], axis=2)

    kernel_stack: AnyNDArray = np.stack([k(X) for k in kernels], axis=2)

    return cast(AnyNDArray, (kernel_stack * starters * stoppers).sum(axis=2))


@pytest.mark.parametrize("variance, lengthscales", [[2.3, 1.4]])
def test_rbf_1d(variance: TensorType, lengthscales: TensorType) -> None:
    X = rng.randn(3, 1)
    kernel = SquaredExponential(lengthscales=lengthscales, variance=variance)

    gram_matrix = kernel(X)
    reference_gram_matrix = ref_rbf_kernel(X, lengthscales, variance)

    assert_allclose(gram_matrix, reference_gram_matrix)


@pytest.mark.parametrize("variance, lengthscales", [[2.3, 1.4]])
def test_rq_1d(variance: TensorType, lengthscales: TensorType) -> None:
    kSE = SquaredExponential(lengthscales=lengthscales, variance=variance)
    kRQ = RationalQuadratic(lengthscales=lengthscales, variance=variance, alpha=1e8)
    rng = np.random.RandomState(1)
    X: AnyNDArray = rng.randn(6, 1).astype(default_float())

    gram_matrix_SE = kSE(X)
    gram_matrix_RQ = kRQ(X)
    assert_allclose(gram_matrix_SE, gram_matrix_RQ)


@check_shapes(
    "variance: []",
    "weight_variances: [broadcast n_active_dims]",
    "bias_variance: []",
    "X: [N, D]",
)
def _assert_arccosine_kern_err(
    variance: TensorType,
    weight_variances: TensorType,
    bias_variance: TensorType,
    order: int,
    X: TensorType,
) -> None:
    kernel = ArcCosine(
        order=order,
        variance=variance,
        weight_variances=weight_variances,
        bias_variance=bias_variance,
    )
    gram_matrix = kernel(X)
    reference_gram_matrix = ref_arccosine_kernel(
        X, order, weight_variances, bias_variance, variance
    )
    assert_allclose(gram_matrix, reference_gram_matrix)


@pytest.mark.parametrize("order", ArcCosine.implemented_orders)
@pytest.mark.parametrize("D, weight_variances", [[1, 1.7], [3, 1.7], [3, (1.1, 1.7, 1.9)]])
@pytest.mark.parametrize("N, bias_variance, variance", [[3, 0.6, 2.3]])
def test_arccosine_1d_and_3d(
    order: int,
    D: int,
    N: int,
    weight_variances: Union[float, Sequence[float]],
    bias_variance: float,
    variance: float,
) -> None:
    X_data = rng.randn(N, D)
    _assert_arccosine_kern_err(variance, weight_variances, bias_variance, order, X_data)


@pytest.mark.parametrize("order", [42])
def test_arccosine_non_implemented_order(order: int) -> None:
    with pytest.raises(ValueError):
        ArcCosine(order=order)


@pytest.mark.parametrize("D, N", [[1, 4]])
def test_arccosine_nan_gradient(D: int, N: int) -> None:
    X = rng.rand(N, D)
    kernel = ArcCosine()
    with tf.GradientTape() as tape:
        Kff = kernel(X)
    grads = tape.gradient(Kff, kernel.trainable_variables)
    assert not np.any(np.isnan(grads))


@pytest.mark.parametrize(
    "base_class",
    [
        SquaredExponential,
        Matern12,
        Matern32,
        Matern52,
    ],
)
@pytest.mark.parametrize(
    "D, lengthscales, period",
    [
        [1, 2.0, 3.0],  # 1d, single lengthscale, single period
        [2, 11.5, 3.0],  # 2d, single lengthscale, single period
        [2, 11.5, (3.0, 6.0)],  # 2d, single lengthscale, ard period
        [2, (11.5, 12.5), 3.0],  # 2d, ard lengthscales, single period
        [2, (11.5, 12.5), (3.0, 6.0)],  # 2d, ard lengthscales, ard period
    ],
)
@pytest.mark.parametrize(
    "N, variance",
    [
        [3, 2.3],
        [5, 1.3],
    ],
)
def test_periodic(
    base_class: Type[Stationary],
    D: int,
    N: int,
    lengthscales: TensorType,
    variance: TensorType,
    period: TensorType,
) -> None:
    X = rng.randn(N, D) if D == 1 else rng.multivariate_normal(np.zeros(D), np.eye(D), N)

    base_kernel = base_class(lengthscales=lengthscales, variance=variance)
    kernel = Periodic(base_kernel, period=period)
    gram_matrix = kernel(X)
    reference_gram_matrix = ref_periodic_kernel(
        X, base_class.__name__, lengthscales, variance, period
    )

    assert_allclose(gram_matrix, reference_gram_matrix)


@pytest.mark.parametrize(
    "base_class",
    [
        SquaredExponential,
        Matern12,
    ],
)
def test_periodic_diag(base_class: Type[Stationary]) -> None:
    N, D = 5, 3
    X = rng.multivariate_normal(np.zeros(D), np.eye(D), N)
    base_kernel = base_class(lengthscales=2.0, variance=1.0)
    kernel = Periodic(base_kernel, period=6.0)
    assert_allclose(base_kernel(X, full_cov=False), kernel(X, full_cov=False))


def test_periodic_non_stationary_base_kernel() -> None:
    error_msg = r"Periodic requires an IsotropicStationary kernel as the `base_kernel`"
    with pytest.raises(TypeError, match=error_msg):
        Periodic(Linear())


def test_periodic_bad_ard_period() -> None:
    error_msg = r"Size of `active_dims` \[1 2\] does not match size of ard parameter \(3\)"
    base_kernel = RBF(active_dims=[1, 2])
    with pytest.raises(ValueError, match=error_msg):
        Periodic(base_kernel, period=[1.0, 1.0, 1.0])


kernel_setups: Tuple[Kernel, ...] = tuple(
    kernel()
    for kernel in gpflow.ci_utils.subclasses(Stationary)
    if kernel not in (IsotropicStationary, AnisotropicStationary)
) + (
    Constant(),
    Linear(),
    Polynomial(),
    ArcCosine(),
)


@pytest.mark.parametrize("D", [1, 5])
@pytest.mark.parametrize("kernel", kernel_setups)
@pytest.mark.parametrize("N", [10])
def test_kernel_symmetry_1d_and_5d(D: int, kernel: Kernel, N: int) -> None:
    X = rng.randn(N, D)
    errors = kernel(X) - kernel(X, X)
    assert np.allclose(errors, 0)


@pytest.mark.parametrize("N, N2, input_dim, output_dim, rank", [[10, 12, 1, 3, 2]])
def test_coregion_shape(N: int, N2: int, input_dim: int, output_dim: int, rank: int) -> None:
    X = np.random.randint(0, output_dim, (N, input_dim))
    X2 = np.random.randint(0, output_dim, (N2, input_dim))
    kernel = Coregion(output_dim=output_dim, rank=rank)
    kernel.W.assign(rng.randn(output_dim, rank))
    kernel.kappa.assign(np.exp(rng.randn(output_dim, 1).reshape(-1)))

    Kff2 = kernel(X, X2)
    assert Kff2.shape == (10, 12)
    Kff = kernel(X)
    assert Kff.shape == (10, 10)


@pytest.mark.parametrize("N, input_dim, output_dim, rank", [[10, 1, 3, 2]])
def test_coregion_diag(N: int, input_dim: int, output_dim: int, rank: int) -> None:
    X = np.random.randint(0, output_dim, (N, input_dim))
    kernel = Coregion(output_dim=output_dim, rank=rank)
    kernel.W.assign(rng.randn(output_dim, rank))
    kernel.kappa.assign(np.exp(rng.randn(output_dim, 1).reshape(-1)))

    K = kernel(X)
    Kdiag = kernel.K_diag(X)
    assert np.allclose(np.diag(K), Kdiag)


@pytest.mark.parametrize("N, input_dim, output_dim, rank", [[10, 1, 3, 2]])
def test_coregion_slice(N: int, input_dim: int, output_dim: int, rank: int) -> None:
    X = np.random.randint(0, output_dim, (N, input_dim))
    X = np.hstack((X, rng.randn(10, 1)))
    kernel1 = Coregion(output_dim=output_dim, rank=rank, active_dims=[0])
    # compute another kernel with additinoal inputs,
    # make sure out kernel is still okay.
    kernel2 = SquaredExponential(active_dims=[1])
    kernel_prod = kernel1 * kernel2
    K1 = kernel_prod(X)
    K2 = kernel1(X) * kernel2(X)  # slicing happens inside kernel
    assert np.allclose(K1, K2)


_dim = 3
kernel_setups_extended: Tuple[Kernel, ...] = (
    kernel_setups
    + (
        SquaredExponential() + Linear(),
        SquaredExponential() * Linear(),
        SquaredExponential() + Linear(variance=rng.rand(_dim)),
    )
    + tuple(ArcCosine(order=order) for order in ArcCosine.implemented_orders)
)


@pytest.mark.parametrize("kernel", kernel_setups_extended)
@pytest.mark.parametrize("N, dim", [[30, _dim]])
def test_diags(kernel: Kernel, N: int, dim: int) -> None:
    X = np.random.randn(N, dim)
    kernel1 = tf.linalg.diag_part(kernel(X, full_cov=True))
    kernel2 = kernel(X, full_cov=False)
    assert np.allclose(kernel1, kernel2)


def test_conv_diag() -> None:
    kernel = Convolutional(SquaredExponential(), [3, 3], [2, 2])
    X = np.random.randn(3, 9)
    kernel_full = np.diagonal(kernel(X, full_cov=True))
    kernel_diag = kernel(X, full_cov=False)
    assert np.allclose(kernel_full, kernel_diag)
    assert 4 == kernel.patch_len
    assert 4 == kernel.num_patches


# Add a rbf and linear kernel, make sure the result is the same as adding the result of
# the kernels separately.
_kernel_setups_add: Tuple[Kernel, ...] = (
    SquaredExponential(),
    Linear(),
    SquaredExponential() + Linear(),
)


@pytest.mark.parametrize("N, D", [[10, 1]])
def test_add_symmetric(N: int, D: int) -> None:
    X = rng.randn(N, D)
    Kffs = [kernel(X) for kernel in _kernel_setups_add]

    assert np.allclose(Kffs[0] + Kffs[1], Kffs[2])


@pytest.mark.parametrize("N, M, D", [[10, 12, 1]])
def test_add_asymmetric(N: int, M: int, D: int) -> None:
    X, Z = rng.randn(N, D), rng.randn(M, D)
    Kfus = [kernel(X, Z) for kernel in _kernel_setups_add]

    assert np.allclose(Kfus[0] + Kfus[1], Kfus[2])


@pytest.mark.parametrize("N, D", [[10, 1]])
def test_white(N: int, D: int) -> None:
    """
    The white kernel should not give the same result when called with k(X) and
    k(X, X)
    """
    X = rng.randn(N, D)
    kernel = White()
    Kff_sym = kernel(X)
    Kff_asym = kernel(X, X)

    assert not np.allclose(Kff_sym, Kff_asym)


_kernel_classes_slice = [
    kernel
    for kernel in gpflow.ci_utils.subclasses(Stationary)
    if kernel not in (IsotropicStationary, AnisotropicStationary)
] + [
    Constant,
    Linear,
    Polynomial,
]

_kernel_triples_slice = [
    (k1(active_dims=[0]), k2(active_dims=[1]), k3(active_dims=slice(0, 1)))
    for k1, k2, k3 in zip(_kernel_classes_slice, _kernel_classes_slice, _kernel_classes_slice)
]


@pytest.mark.parametrize("kernel_triple", _kernel_triples_slice)
@pytest.mark.parametrize("N, D", [[20, 2]])
def test_slice_symmetric(kernel_triple: Tuple[Kernel, Kernel, Kernel], N: int, D: int) -> None:
    X = rng.randn(N, D)
    K1, K3 = kernel_triple[0](X), kernel_triple[2](X[:, :1])
    assert np.allclose(K1, K3)
    K2, K4 = kernel_triple[1](X), kernel_triple[2](X[:, 1:])
    assert np.allclose(K2, K4)


@pytest.mark.parametrize("kernel_triple", _kernel_triples_slice)
@pytest.mark.parametrize("N, M, D", [[10, 12, 2]])
def test_slice_asymmetric(
    kernel_triple: Tuple[Kernel, Kernel, Kernel], N: int, M: int, D: int
) -> None:
    X = rng.randn(N, D)
    Z = rng.randn(M, D)
    K1, K3 = kernel_triple[0](X, Z), kernel_triple[2](X[:, :1], Z[:, :1])
    assert np.allclose(K1, K3)
    K2, K4 = kernel_triple[1](X, Z), kernel_triple[2](X[:, 1:], Z[:, 1:])
    assert np.allclose(K2, K4)


_kernel_setups_prod: Tuple[Kernel, ...] = (
    Matern32(),
    Matern52(lengthscales=0.3),
    Matern32() * Matern52(lengthscales=0.3),
)


@pytest.mark.parametrize("N, D", [[30, 2]])
def test_product(N: int, D: int) -> None:
    X = rng.randn(N, D)
    Kffs = [kernel(X) for kernel in _kernel_setups_prod]

    assert np.allclose(Kffs[0] * Kffs[1], Kffs[2])


@pytest.mark.parametrize("N, D", [[30, 4], [10, 7]])
def test_active_product(N: int, D: int) -> None:
    X = rng.randn(N, D)
    dims, rand_idx, ls = (
        list(range(D)),
        int(rng.randint(0, D)),
        rng.uniform(1.0, 7.0, D),
    )
    active_dims_list = [dims[:rand_idx] + dims[rand_idx + 1 :], [rand_idx], dims]
    lengthscales_list = [
        np.hstack([ls[:rand_idx], ls[rand_idx + 1 :]]),
        ls[rand_idx],
        ls,
    ]
    kernels = [
        SquaredExponential(lengthscales=lengthscales, active_dims=dims)
        for dims, lengthscales in zip(active_dims_list, lengthscales_list)
    ]
    kernel_prod = kernels[0] * kernels[1]

    Kff = kernels[2](X)
    Kff_prod = kernel_prod(X)

    assert np.allclose(Kff, Kff_prod)


@pytest.mark.parametrize("D", [4, 7])
def test_ard_init_scalar(D: int) -> None:
    """
    For ard kernels, make sure that kernels can be instantiated with a single
    scalar lengthscale or a suitable array of lengthscales
    """
    kernel_1 = SquaredExponential(lengthscales=2.3)
    kernel_2 = SquaredExponential(lengthscales=np.ones(D) * 2.3)
    lengthscales_1 = kernel_1.lengthscales.numpy()
    lengthscales_2 = kernel_2.lengthscales.numpy()
    assert np.allclose(lengthscales_1, lengthscales_2, atol=1e-10)


def test_ard_invalid_active_dims() -> None:
    msg = r"Size of `active_dims` \[1\] does not match size of ard parameter \(2\)"
    with pytest.raises(ValueError, match=msg):
        SquaredExponential(lengthscales=np.ones(2), active_dims=[1])


@pytest.mark.parametrize(
    "kernel_class, param_name",
    [
        [SquaredExponential, "lengthscales"],
        [Linear, "variance"],
        [ArcCosine, "weight_variances"],
        [Cosine, "lengthscales"],
    ],
)
@pytest.mark.parametrize(
    "param_value, ard",
    [
        [1.0, False],
        [[1.0], True],
        [[1.0, 1.0], True],
    ],
)
def test_ard_property(
    kernel_class: Type[Kernel], param_name: str, param_value: Any, ard: bool
) -> None:
    kernel = kernel_class(**{param_name: param_value})
    assert kernel.ard is ard


@pytest.mark.parametrize(
    "locations, steepness, error_msg",
    [
        # 1. Kernels locations dimension mismatch
        [
            [1.0],
            1.0,
            r"Number of kernels \(3\) must be one more than the number of changepoint locations \(1\)",
        ],
        # 2. Locations steepness dimension mismatch
        [
            [1.0, 2.0],
            [1.0],
            r"Dimension of steepness \(1\) does not match number of changepoint locations \(2\)",
        ],
    ],
)
def test_changepoints_init_fail(
    locations: Sequence[float], steepness: Union[float, Sequence[float]], error_msg: str
) -> None:
    kernels: Sequence[Kernel] = [
        Matern12(),
        Linear(),
        Matern32(),
    ]
    with pytest.raises(ValueError, match=error_msg):
        ChangePoints(kernels, locations, steepness)


@check_shapes(
    "X: [N, D]",
    "locations: [L]",
    "steepness: [broadcast L]",
)
def _assert_changepoints_kern_err(
    X: TensorType,
    kernels: Sequence[Kernel],
    locations: Sequence[float],
    steepness: Union[float, Sequence[float]],
) -> None:
    kernel = ChangePoints(kernels, locations, steepness=steepness)
    reference_gram_matrix = _ref_changepoints(X, kernels, locations, steepness)

    assert_allclose(kernel(X), reference_gram_matrix)
    assert_allclose(kernel.K_diag(X), np.diag(reference_gram_matrix))


@pytest.mark.parametrize("N", [2, 10])
@pytest.mark.parametrize(
    "kernels, locations, steepness",
    [
        # 1. Single changepoint
        [[Constant(), Constant()], [2.0], 5.0],
        # 2. Two changepoints
        [
            [
                Constant(),
                Constant(),
                Constant(),
            ],
            [1.0, 2.0],
            5.0,
        ],
        # 3. Multiple steepness
        [
            [
                Constant(),
                Constant(),
                Constant(),
            ],
            [1.0, 2.0],
            [5.0, 10.0],
        ],
        # 4. Variety of kernels
        [
            [
                Matern12(),
                Linear(),
                SquaredExponential(),
                Constant(),
            ],
            [1.0, 2.0, 3.0],
            5.0,
        ],
    ],
)
def test_changepoints(
    N: int,
    kernels: Sequence[Kernel],
    locations: Sequence[float],
    steepness: Union[float, Sequence[float]],
) -> None:
    X_data = rng.randn(N, 1)
    _assert_changepoints_kern_err(X_data, kernels, locations, steepness)


@pytest.mark.parametrize(
    "active_dims_1, active_dims_2, is_separate",
    [
        [[1, 2, 3], None, False],
        [None, [1, 2, 3], False],
        [None, None, False],
        [[1, 2, 3], [3, 4, 5], False],
        [[1, 2, 3], [4, 5, 6], True],
    ],
)
def test_on_separate_dims(
    active_dims_1: Optional[Sequence[int]],
    active_dims_2: Optional[Sequence[int]],
    is_separate: bool,
) -> None:
    kernel_1 = Linear(active_dims=active_dims_1)
    kernel_2 = SquaredExponential(active_dims=active_dims_2)
    assert kernel_1.on_separate_dims(kernel_2) == is_separate
    assert kernel_2.on_separate_dims(kernel_1) == is_separate
    assert kernel_1.on_separate_dims(kernel_1) is False
    assert kernel_2.on_separate_dims(kernel_2) is False


@pytest.mark.parametrize("kernel", kernel_setups_extended)
def test_kernel_call_diag_and_X2_errors(kernel: Kernel) -> None:
    X = rng.randn(4, 1)
    X2 = rng.randn(5, 1)

    with pytest.raises(ValueError):
        kernel(X, X2, full_cov=False)


def test_periodic_active_dims_matches() -> None:
    active_dims = [1]
    base_kernel = SquaredExponential(active_dims=active_dims)
    kernel = Periodic(base_kernel=base_kernel)

    assert kernel.active_dims == base_kernel.active_dims

    # type-ignores below, is because mypy doesn't understand that the setter and the getter for
    # `active_dims` have different types.

    kernel.active_dims = [2]  # type: ignore[assignment]
    assert kernel.active_dims == base_kernel.active_dims

    base_kernel.active_dims = [3]  # type: ignore[assignment]
    assert kernel.active_dims == base_kernel.active_dims


def test_latent_kernels() -> None:
    kernel_list: Tuple[Kernel, ...] = (SquaredExponential(), White(), White() + Linear())

    multioutput_kernel_list: Tuple[MultioutputKernel, ...] = (
        SharedIndependent(SquaredExponential(), 3),
        SeparateIndependent(kernel_list),
        LinearCoregionalization(kernel_list, np.random.random((5, 3))),
    )
    assert len(multioutput_kernel_list[0].latent_kernels) == 1
    assert multioutput_kernel_list[1].latent_kernels == tuple(kernel_list)
    assert multioutput_kernel_list[2].latent_kernels == tuple(kernel_list)


def test_combination_LMC_kernels() -> None:
    N, D, P = 100, 3, 2
    kernel_list1: Tuple[Kernel, ...] = (Linear(active_dims=[1]), SquaredExponential())
    L1 = len(kernel_list1)
    kernel_list2: Tuple[Kernel, ...] = (SquaredExponential(), Linear(), Linear())
    L2 = len(kernel_list2)
    k1 = LinearCoregionalization(kernel_list1, np.random.randn(P, L1))
    k2 = LinearCoregionalization(kernel_list2, np.random.randn(P, L2))
    kernel = k1 + k2
    X = np.random.randn(N, D)
    K1 = k1(X, full_cov=True)
    K2 = k2(X, full_cov=True)
    K = kernel(X, full_cov=True)
    assert K.shape == [N, P, N, P]
    np.testing.assert_allclose(K, K1 + K2)


# tests/gpflow/kernels/test_positive_semidefinite.py

# Copyright 2018 the GPflow authors.
#
# 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 typing import Type

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_array_less

import gpflow.ci_utils
from gpflow import kernels

KERNEL_CLASSES = [
    kernel
    for cls in (kernels.Static, kernels.Stationary, kernels.Linear)
    for kernel in gpflow.ci_utils.subclasses(cls)
    if kernel not in (kernels.IsotropicStationary, kernels.AnisotropicStationary)
] + [kernels.ArcCosine]

rng = np.random.RandomState(42)


def pos_semidefinite(kernel: kernels.Kernel) -> None:
    N, D = 100, 5
    X = rng.randn(N, D)

    cov = kernel(X)
    eig = tf.linalg.eigvalsh(cov).numpy()
    assert_array_less(-1e-12, eig)


@pytest.mark.parametrize("kernel_class", KERNEL_CLASSES)
def test_positive_semidefinite(kernel_class: Type[kernels.Kernel]) -> None:
    """
    A valid kernel is positive semidefinite. Some kernels are only valid for
    particular input shapes, see https://github.com/GPflow/GPflow/issues/1328
    """
    kernel = kernel_class()
    pos_semidefinite(kernel)


@pytest.mark.parametrize(
    "base_class", [kernel for kernel in gpflow.ci_utils.subclasses(kernels.IsotropicStationary)]
)
def test_positive_semidefinite_periodic(base_class: Type[kernels.IsotropicStationary]) -> None:
    """
    A valid kernel is positive semidefinite. Some kernels are only valid for
    particular input shapes, see https://github.com/GPflow/GPflow/issues/1328
    """
    kernel = kernels.Periodic(base_class())
    pos_semidefinite(kernel)


# tests/gpflow/kernels/test_scaled_euclid_dist.py

# Copyright 2018 the GPflow authors.
#
# 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.

import numpy as np
import pytest
import tensorflow as tf

import gpflow.kernels as kernels
from gpflow.base import AnyNDArray

rng = np.random.RandomState(0)


class Datum:
    num_data = 100
    D = 100
    X: AnyNDArray = rng.rand(num_data, D) * 100


kernel_list = [
    kernels.Matern12(),
    kernels.Matern32(),
    kernels.Matern52(),
    kernels.Exponential(),
    kernels.Cosine(),
]


@pytest.mark.parametrize("kernel", kernel_list)
def test_kernel_euclidean_distance(kernel: kernels.Kernel) -> None:
    """
    Tests output & gradients of kernels that are a function of the (scaled) euclidean distance
    of the points. We test on a high dimensional space, which can generate very small distances
    causing the scaled_square_dist to generate some negative values.
    """
    K = kernel(Datum.X)
    assert not np.isnan(K).any(), "NaNs in the output of the " + kernel.__name__ + "kernel."
    assert np.isfinite(K).all(), "Infs in the output of the " + kernel.__name__ + " kernel."

    X_as_param = tf.Variable(Datum.X)
    with tf.GradientTape() as tape:
        K_value = kernel(X_as_param, X_as_param)
    dK = tape.gradient(K_value, X_as_param)[0]

    assert not np.isnan(dK).any(), "NaNs in the gradient of the " + kernel.__name__ + " kernel."
    assert np.isfinite(dK).all(), "Infs in the output of the " + kernel.__name__ + " kernel."


# tests/gpflow/likelihoods/test_function_params.py

# Copyright 2022 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 dataclasses import dataclass
from typing import Any

import numpy as np
import pytest
import tensorflow as tf

from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.functions import Linear
from gpflow.likelihoods import Beta, Gamma, Gaussian, ScalarLikelihood, StudentT
from gpflow.utilities import to_default_float


class Datum:
    cs = ShapeChecker().check_shapes
    rng = np.random.default_rng(20220623)

    batch = (1, 2)
    N = 5
    D = 3
    Q = 2
    X_shape = (*batch, N, D)
    Y_shape = (*batch, N, Q)

    np_X_positive = rng.random(X_shape)
    np_X_positive[:, :, :, 0] = np.linspace(0.1, 1.0, N)[None, None, :]
    X_positive = to_default_float(np_X_positive)
    X_negative = to_default_float(-rng.random(X_shape))
    F = to_default_float(0.5 * np.ones(Y_shape))
    Fmu = to_default_float(0.5 * np.ones(Y_shape))
    Fvar = to_default_float(0.1 * np.ones(Y_shape))
    Y = to_default_float(0.5 * np.ones(Y_shape))

    linear = Linear(A=[[1.0, 1.0], [0.0, 0.0], [0.0, 0.0]], b=0.0)


def diff(x: tf.Tensor, *, axis: int) -> tf.Tensor:
    s = [slice(None, None) for _ in x.shape]
    s[axis] = slice(None, -1)
    head = x.__getitem__(s)
    s[axis] = slice(1, None)
    tail = x.__getitem__(s)
    return tail - head


def assert_decreasing(x: tf.Tensor, *, axis: int) -> None:
    decreasing = diff(x, axis=axis) < 0
    assert tf.reduce_all(decreasing), x


def assert_increasing(x: tf.Tensor, *, axis: int) -> None:
    increasing = diff(x, axis=axis) > 0
    assert tf.reduce_all(increasing), x


def assert_constant(x: tf.Tensor, *, axis: int) -> None:
    same = diff(x, axis=axis) == 0
    assert tf.reduce_all(same), x


def no_assert(x: tf.Tensor, *, axis: int) -> None:
    pass


@dataclass
class LikelihoodSetup:
    name: str
    likelihood: ScalarLikelihood
    likelihood_assert: Any
    mean_assert: Any
    variance_assert: Any
    variational_expectations_assert: Any

    @property
    def __name__(self) -> str:  # Used by pytest to generate (pretty) test ids.
        return self.name


LIKELIHOODS = [
    LikelihoodSetup(
        name="gaussian_variance",
        likelihood=Gaussian(variance=Datum.linear),
        likelihood_assert=assert_decreasing,
        mean_assert=assert_constant,
        variance_assert=assert_increasing,
        variational_expectations_assert=assert_decreasing,
    ),
    LikelihoodSetup(
        name="gaussian_scale",
        likelihood=Gaussian(scale=Datum.linear),
        likelihood_assert=assert_decreasing,
        mean_assert=assert_constant,
        variance_assert=assert_increasing,
        variational_expectations_assert=no_assert,
    ),
    LikelihoodSetup(
        name="student_t",
        likelihood=StudentT(scale=Datum.linear),
        likelihood_assert=assert_decreasing,
        mean_assert=assert_constant,
        variance_assert=assert_increasing,
        variational_expectations_assert=no_assert,
    ),
    LikelihoodSetup(
        name="gamma",
        likelihood=Gamma(shape=Datum.linear),
        likelihood_assert=no_assert,
        mean_assert=assert_increasing,
        variance_assert=assert_increasing,
        variational_expectations_assert=no_assert,
    ),
    LikelihoodSetup(
        name="beta",
        likelihood=Beta(scale=Datum.linear),
        likelihood_assert=no_assert,
        mean_assert=assert_constant,
        variance_assert=assert_decreasing,
        variational_expectations_assert=assert_increasing,
    ),
]


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_log_prob__positive(setup: LikelihoodSetup) -> None:
    lp = setup.likelihood.log_prob(Datum.X_positive, Datum.F, Datum.Y)
    setup.likelihood_assert(lp, axis=-1)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_log_prob__negative(setup: LikelihoodSetup) -> None:
    # We expect the negative parameter values to be clamped to the same value, so output should be
    # constant:
    lp = setup.likelihood.log_prob(Datum.X_negative, Datum.F, Datum.Y)
    assert_constant(lp, axis=-1)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_conditional_mean__positive(setup: LikelihoodSetup) -> None:
    # We don't expect X to affect the mean.
    cm = setup.likelihood.conditional_mean(Datum.X_positive, Datum.F)
    setup.mean_assert(cm, axis=-2)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_conditional_mean__negative(setup: LikelihoodSetup) -> None:
    # We expect the negative parameter values to be clamped to the same value, so output should be
    # constant:
    cm = setup.likelihood.conditional_mean(Datum.X_negative, Datum.F)
    assert_constant(cm, axis=-2)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_conditional_variance__positive(setup: LikelihoodSetup) -> None:
    cv = setup.likelihood.conditional_variance(Datum.X_positive, Datum.F)
    setup.variance_assert(cv, axis=-2)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_conditional_variance__negative(setup: LikelihoodSetup) -> None:
    # We expect the negative parameter values to be clamped to the same value, so output should be
    # constant:
    cv = setup.likelihood.conditional_variance(Datum.X_negative, Datum.F)
    assert_constant(cv, axis=-2)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_predict_mean_and_var__positive(setup: LikelihoodSetup) -> None:
    mu, var = setup.likelihood.predict_mean_and_var(Datum.X_positive, Datum.Fmu, Datum.Fvar)
    setup.mean_assert(mu, axis=-2)
    setup.variance_assert(var, axis=-2)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_predict_mean_and_var__negative(setup: LikelihoodSetup) -> None:
    # We expect the negative parameter values to be clamped to the same value, so output should be
    # constant:
    mu, var = setup.likelihood.predict_mean_and_var(Datum.X_negative, Datum.Fmu, Datum.Fvar)
    assert_constant(mu, axis=-2)
    assert_constant(var, axis=-2)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_predict_log_density__positive(setup: LikelihoodSetup) -> None:
    ld = setup.likelihood.predict_log_density(Datum.X_positive, Datum.Fmu, Datum.Fvar, Datum.Y)
    setup.likelihood_assert(ld, axis=-1)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_predict_log_density__negative(setup: LikelihoodSetup) -> None:
    # We expect the negative parameter values to be clamped to the same value, so output should be
    # constant:
    ld = setup.likelihood.predict_log_density(Datum.X_negative, Datum.Fmu, Datum.Fvar, Datum.Y)
    assert_constant(ld, axis=-2)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_variational_expectation__positive(setup: LikelihoodSetup) -> None:
    ve = setup.likelihood.variational_expectations(Datum.X_positive, Datum.Fmu, Datum.Fvar, Datum.Y)
    setup.variational_expectations_assert(ve, axis=-1)


@pytest.mark.parametrize("setup", LIKELIHOODS)
def test_variational_expectation__negative(setup: LikelihoodSetup) -> None:
    # We expect the negative parameter values to be clamped to the same value, so output should be
    # constant:
    ve = setup.likelihood.variational_expectations(Datum.X_negative, Datum.Fmu, Datum.Fvar, Datum.Y)
    assert_constant(ve, axis=-2)


# tests/gpflow/likelihoods/test_heteroskedastic.py

# Copyright 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.

import numpy as np
import tensorflow as tf

from gpflow.base import AnyNDArray
from gpflow.likelihoods import HeteroskedasticTFPConditional

tf.random.set_seed(99012)


class Data:
    rng = np.random.RandomState(123)
    N = 5
    X = np.linspace(0, 1, num=N)[:, None]
    Y = rng.randn(N, 1)
    f_mean = rng.randn(N, 2)
    f_var: AnyNDArray = rng.randn(N, 2) ** 2


def test_analytic_mean_and_var() -> None:
    """
    Test that quadrature computation used in HeteroskedasticTFPConditional
    of the predictive mean and variance is close to the analytical version,
    which can be computed for the special case of N(y | mean=f1, scale=exp(f2)),
    where f1, f2 ~ GP.
    """
    analytic_mean = Data.f_mean[:, [0]]
    analytic_variance = np.exp(Data.f_mean[:, [1]] + Data.f_var[:, [1]]) ** 2 + Data.f_var[:, [0]]

    likelihood = HeteroskedasticTFPConditional()
    y_mean, y_var = likelihood.predict_mean_and_var(Data.X, Data.f_mean, Data.f_var)

    np.testing.assert_allclose(y_mean, analytic_mean)
    np.testing.assert_allclose(y_var, analytic_variance, rtol=1.5e-6)


# tests/gpflow/likelihoods/test_heteroskedastic_constant_variance.py

# Copyright 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 typing import Tuple, Type

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp
from _pytest.fixtures import SubRequest

import gpflow
from gpflow.base import AnyNDArray, TensorType
from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.likelihoods import HeteroskedasticTFPConditional

tf.random.set_seed(99012)


EquivalentLikelihoods = Tuple[
    gpflow.likelihoods.ScalarLikelihood, gpflow.likelihoods.HeteroskedasticTFPConditional
]


class Data:
    cs = ShapeChecker().check_shape

    g_var = 0.345
    rng = np.random.RandomState(123)
    N = 5
    X = cs(rng.randn(N, 2), "[N, D]")
    Y = cs(rng.randn(N, 1), "[N, P]")
    # single "GP" (for the mean):
    f_mean = cs(rng.randn(N, 1), "[N, Q]")
    f_var: AnyNDArray = cs(rng.randn(N, 1) ** 2, "[N, Q]")  # ensure positivity
    equivalent_f2 = cs(np.log(g_var) / 2, "[]")
    f2_mean = cs(np.full((N, 1), equivalent_f2), "[N, Q]")
    f2_var = cs(np.zeros((N, 1)), "[N, Q]")
    F2_mean = cs(np.c_[f_mean, f2_mean], "[N, Q2]")
    F2_var = cs(np.c_[f_var, f2_var], "[N, Q2]")


def student_t_class_factory(df: int = 3) -> Type[tfp.distributions.StudentT]:
    r"""
    Returns tfp.Distribution.StudentT class (not instance!)
    where df (degrees of freedom) is pre-specified.

    This class allows to instantiate a StundentT object by passing
    loc and sale at initialisation for a given degree-of-freedom.
    """

    class _StudentT(tfp.distributions.StudentT):
        def __init__(self, loc: TensorType, scale: TensorType) -> None:
            super().__init__(df, loc=loc, scale=scale)

    return _StudentT


@pytest.fixture(name="equivalent_likelihoods", params=["studentt", "gaussian"])
def _equivalent_likelihoods_fixture(
    request: SubRequest,
) -> EquivalentLikelihoods:
    if request.param == "studentt":
        return (
            gpflow.likelihoods.StudentT(scale=Data.g_var ** 0.5, df=3.0),
            HeteroskedasticTFPConditional(distribution_class=student_t_class_factory(df=3)),
        )
    elif request.param == "gaussian":
        return (
            gpflow.likelihoods.Gaussian(variance=Data.g_var),
            HeteroskedasticTFPConditional(distribution_class=tfp.distributions.Normal),
        )
    assert False, f"Unknown likelihood {request.param}."


def test_log_prob(equivalent_likelihoods: EquivalentLikelihoods) -> None:
    """
    heteroskedastic likelihood where the variance parameter is always constant
     giving the same answers for variational_expectations, predict_mean_and_var,
      etc as the regular Gaussian  likelihood
    """
    homoskedastic_likelihood, heteroskedastic_likelihood = equivalent_likelihoods
    np.testing.assert_array_almost_equal(
        homoskedastic_likelihood.log_prob(Data.X, Data.f_mean, Data.Y),
        heteroskedastic_likelihood.log_prob(Data.X, Data.F2_mean, Data.Y),
    )


def test_variational_expectations(equivalent_likelihoods: EquivalentLikelihoods) -> None:
    homoskedastic_likelihood, heteroskedastic_likelihood = equivalent_likelihoods
    np.testing.assert_array_almost_equal(
        homoskedastic_likelihood.variational_expectations(Data.X, Data.f_mean, Data.f_var, Data.Y),
        heteroskedastic_likelihood.variational_expectations(
            Data.X, Data.F2_mean, Data.F2_var, Data.Y
        ),
        decimal=2,  # student-t case has a max absolute difference of 0.0034
    )


def test_predict_mean_and_var(equivalent_likelihoods: EquivalentLikelihoods) -> None:
    homoskedastic_likelihood, heteroskedastic_likelihood = equivalent_likelihoods
    np.testing.assert_allclose(
        homoskedastic_likelihood.predict_mean_and_var(Data.X, Data.f_mean, Data.f_var),
        heteroskedastic_likelihood.predict_mean_and_var(Data.X, Data.F2_mean, Data.F2_var),
    )


def test_conditional_mean(equivalent_likelihoods: EquivalentLikelihoods) -> None:
    homoskedastic_likelihood, heteroskedastic_likelihood = equivalent_likelihoods
    np.testing.assert_allclose(
        homoskedastic_likelihood.conditional_mean(Data.X, Data.f_mean),
        heteroskedastic_likelihood.conditional_mean(Data.X, Data.F2_mean),
    )


def test_conditional_variance(equivalent_likelihoods: EquivalentLikelihoods) -> None:
    homoskedastic_likelihood, heteroskedastic_likelihood = equivalent_likelihoods
    np.testing.assert_allclose(
        homoskedastic_likelihood.conditional_variance(Data.X, Data.f_mean),
        heteroskedastic_likelihood.conditional_variance(Data.X, Data.F2_mean),
    )


def test_predict_log_density(equivalent_likelihoods: EquivalentLikelihoods) -> None:
    homoskedastic_likelihood, heteroskedastic_likelihood = equivalent_likelihoods
    np.testing.assert_array_almost_equal(
        homoskedastic_likelihood.predict_log_density(Data.X, Data.f_mean, Data.f_var, Data.Y),
        heteroskedastic_likelihood.predict_log_density(Data.X, Data.F2_mean, Data.F2_var, Data.Y),
        decimal=1,  # student-t has a max absolute difference of 0.025
    )


# tests/gpflow/likelihoods/test_likelihoods.py

# Copyright 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 typing import Sequence, Tuple

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
import gpflow.ci_utils
from gpflow.base import TensorType
from gpflow.config import default_float, default_int

from gpflow.likelihoods import (  # isort:skip
    # classes we cannot test:
    HeteroskedasticTFPConditional,
    Likelihood,
    MonteCarloLikelihood,
    MultiLatentLikelihood,
    MultiLatentTFPConditional,
    QuadratureLikelihood,
    ScalarLikelihood,
    # classes we do test in this file:
    Bernoulli,
    Beta,
    Exponential,
    Gamma,
    Gaussian,
    GaussianMC,
    MultiClass,
    Ordinal,
    Poisson,
    StudentT,
)

tf.random.set_seed(99012)


class Datum:
    tolerance = 1e-06
    N = 10
    Xshape = (N, 2)
    Yshape = (N, 3)
    X = tf.random.normal(Xshape, dtype=tf.float64)
    Y = tf.random.normal(Yshape, dtype=tf.float64)
    F = tf.random.normal(Yshape, dtype=tf.float64)
    Fmu = tf.random.normal(Yshape, dtype=tf.float64)
    Fvar = 0.01 * tf.random.normal(Yshape, dtype=tf.float64) ** 2
    Fvar_zero = tf.zeros(Yshape, dtype=tf.float64)


class LikelihoodSetup(object):
    def __init__(
        self,
        likelihood: Likelihood,
        Y: TensorType = Datum.Y,
        rtol: float = 1e-06,
        atol: float = 0.0,
    ) -> None:
        self.likelihood = likelihood
        self.Y = Y
        self.rtol = rtol
        self.atol = atol

    def __repr__(self) -> str:
        name = self.likelihood.__class__.__name__
        return f"{name}-rtol={self.rtol}-atol={self.atol}"


scalar_likelihood_setups = [
    LikelihoodSetup(Gaussian()),
    LikelihoodSetup(StudentT()),
    LikelihoodSetup(Beta(), Y=tf.random.uniform(Datum.Yshape, dtype=default_float())),
    LikelihoodSetup(
        Ordinal(np.array([-1, 1])),
        Y=tf.random.uniform(Datum.Yshape, 0, 3, dtype=default_int()),
    ),
    LikelihoodSetup(
        Poisson(invlink=tf.square),
        Y=tf.random.poisson(Datum.Yshape, 1.0, dtype=default_float()),
    ),
    LikelihoodSetup(
        Exponential(invlink=tf.square),
        Y=tf.random.uniform(Datum.Yshape, dtype=default_float()),
    ),
    LikelihoodSetup(
        Gamma(invlink=tf.square),
        Y=tf.random.uniform(Datum.Yshape, dtype=default_float()),
    ),
    LikelihoodSetup(
        Bernoulli(invlink=tf.sigmoid),
        Y=tf.random.uniform(Datum.Yshape, dtype=default_float()),
    ),
]

likelihood_setups = scalar_likelihood_setups + [
    LikelihoodSetup(
        MultiClass(3),
        Y=tf.argmax(Datum.Y, 1).numpy().reshape(-1, 1),
        rtol=1e-3,
        atol=1e-3,
    ),
]


def filter_analytic_scalar_likelihood(method_name: str) -> Sequence[LikelihoodSetup]:
    assert method_name in (
        "_variational_expectations",
        "_predict_log_density",
        "_predict_mean_and_var",
    )

    def is_analytic(likelihood: Likelihood) -> bool:
        assert not isinstance(likelihood, MonteCarloLikelihood)
        assert isinstance(likelihood, ScalarLikelihood)
        quadrature_fallback = getattr(ScalarLikelihood, method_name)
        actual_method = getattr(likelihood.__class__, method_name)
        return actual_method is not quadrature_fallback

    return [l for l in scalar_likelihood_setups if is_analytic(get_likelihood(l))]


def get_likelihood(likelihood_setup: LikelihoodSetup) -> Likelihood:
    if isinstance(likelihood_setup, type(pytest.param())):
        (likelihood_setup,) = likelihood_setup.values
    return likelihood_setup.likelihood


def test_no_missing_likelihoods() -> None:
    tested_likelihood_types = [get_likelihood(l).__class__ for l in likelihood_setups]
    for likelihood_class in gpflow.ci_utils.subclasses(Likelihood):
        if likelihood_class in (
            QuadratureLikelihood,
            ScalarLikelihood,
            MonteCarloLikelihood,
            MultiLatentLikelihood,
        ):
            # abstract base classes that cannot be tested
            continue

        if likelihood_class in tested_likelihood_types:
            # tested by parametrized tests (see test_multiclass.py for MultiClass quadrature)
            continue

        if likelihood_class is gpflow.likelihoods.SwitchedLikelihood:
            # tested separately, see test_switched_likelihood.py
            continue

        if likelihood_class in (MultiLatentTFPConditional, HeteroskedasticTFPConditional):
            # tested separately, see test_heteroskedastic*.py
            continue

        if issubclass(likelihood_class, MonteCarloLikelihood):
            if likelihood_class is GaussianMC:
                continue  # tested explicitly by test_montecarlo_*
            if likelihood_class is gpflow.likelihoods.Softmax:
                continue  # tested explicitly by test_softmax_{y_shape_assert,bernoulli_equivalence}, see test_multiclass.py

        assert False, f"no test for likelihood class {likelihood_class}"


@pytest.mark.parametrize("likelihood_setup", likelihood_setups)
@pytest.mark.parametrize("X, mu, var", [[Datum.X, Datum.Fmu, tf.zeros_like(Datum.Fmu)]])
def test_conditional_mean_and_variance(
    likelihood_setup: LikelihoodSetup, X: TensorType, mu: TensorType, var: TensorType
) -> None:
    """
    Here we make sure that the conditional_mean and conditional_var functions
    give the same result as the predict_mean_and_var function if the prediction
    has no uncertainty.
    """
    mu1 = likelihood_setup.likelihood.conditional_mean(X, mu)
    var1 = likelihood_setup.likelihood.conditional_variance(X, mu)
    mu2, var2 = likelihood_setup.likelihood.predict_mean_and_var(X, mu, var)
    assert_allclose(mu1, mu2, rtol=likelihood_setup.rtol, atol=likelihood_setup.atol)
    assert_allclose(var1, var2, rtol=likelihood_setup.rtol, atol=likelihood_setup.atol)


@pytest.mark.parametrize("likelihood_setup", likelihood_setups)
def test_variational_expectations(likelihood_setup: LikelihoodSetup) -> None:
    """
    Here we make sure that the variational_expectations gives the same result
    as log_prob if the latent function has no uncertainty.
    """
    likelihood = likelihood_setup.likelihood
    X = Datum.X
    F = Datum.F
    Y = likelihood_setup.Y
    r1 = likelihood.log_prob(X, F, Y)
    r2 = likelihood.variational_expectations(X, F, tf.zeros_like(F), Y)
    assert_allclose(r1, r2, atol=likelihood_setup.atol, rtol=likelihood_setup.rtol)


@pytest.mark.parametrize(
    "likelihood_setup", filter_analytic_scalar_likelihood("_variational_expectations")
)
@pytest.mark.parametrize("mu, var", [[Datum.Fmu, Datum.Fvar]])
def test_scalar_likelihood_quadrature_variational_expectation(
    likelihood_setup: LikelihoodSetup, mu: TensorType, var: TensorType
) -> None:
    """
    Where quadrature methods have been overwritten, make sure the new code
    does something close to the quadrature.
    """
    x = Datum.X
    likelihood, y = likelihood_setup.likelihood, likelihood_setup.Y
    F1 = likelihood.variational_expectations(x, mu, var, y)
    F2 = ScalarLikelihood.variational_expectations(likelihood, x, mu, var, y)
    assert_allclose(F1, F2, rtol=likelihood_setup.rtol, atol=likelihood_setup.atol)


@pytest.mark.parametrize(
    "likelihood_setup", filter_analytic_scalar_likelihood("_predict_log_density")
)
@pytest.mark.parametrize("mu, var", [[Datum.Fmu, Datum.Fvar]])
def test_scalar_likelihood_quadrature_predict_log_density(
    likelihood_setup: LikelihoodSetup, mu: TensorType, var: TensorType
) -> None:
    x = Datum.X
    likelihood, y = likelihood_setup.likelihood, likelihood_setup.Y
    F1 = likelihood.predict_log_density(x, mu, var, y)
    F2 = ScalarLikelihood.predict_log_density(likelihood, x, mu, var, y)
    assert_allclose(F1, F2, rtol=likelihood_setup.rtol, atol=likelihood_setup.atol)


@pytest.mark.parametrize(
    "likelihood_setup", filter_analytic_scalar_likelihood("_predict_mean_and_var")
)
@pytest.mark.parametrize("X, mu, var", [[Datum.X, Datum.Fmu, Datum.Fvar]])
def test_scalar_likelihood_quadrature_predict_mean_and_var(
    likelihood_setup: LikelihoodSetup, X: TensorType, mu: TensorType, var: TensorType
) -> None:
    likelihood = likelihood_setup.likelihood
    F1m, F1v = likelihood.predict_mean_and_var(X, mu, var)
    F2m, F2v = ScalarLikelihood.predict_mean_and_var(likelihood, X, mu, var)
    assert_allclose(F1m, F2m, rtol=likelihood_setup.rtol, atol=likelihood_setup.atol)
    assert_allclose(F1v, F2v, rtol=likelihood_setup.rtol, atol=likelihood_setup.atol)


def _make_montecarlo_x_mu_var_y() -> Sequence[tf.Tensor]:
    x_mu_var_y = [tf.random.normal((3, 10), dtype=tf.float64)] * 4
    x_mu_var_y[2] = 0.01 * (x_mu_var_y[2] ** 2)
    return x_mu_var_y


def _make_montecarlo_likelihoods(var: float) -> Tuple[GaussianMC, Gaussian]:
    gaussian_mc_likelihood = GaussianMC(var)
    gaussian_mc_likelihood.num_monte_carlo_points = 1000000
    return gaussian_mc_likelihood, Gaussian(var)


@pytest.mark.parametrize("likelihood_var", [0.3, 0.5, 1])
@pytest.mark.parametrize("x, mu, var, y", [_make_montecarlo_x_mu_var_y()])
def test_montecarlo_variational_expectation(
    likelihood_var: float, x: TensorType, mu: TensorType, var: TensorType, y: TensorType
) -> None:
    likelihood_gaussian_mc, likelihood_gaussian = _make_montecarlo_likelihoods(likelihood_var)
    assert_allclose(
        likelihood_gaussian_mc.variational_expectations(x, mu, var, y),
        likelihood_gaussian.variational_expectations(x, mu, var, y),
        rtol=5e-4,
        atol=1e-4,
    )


@pytest.mark.parametrize("likelihood_var", [0.3, 0.5, 1.0])
@pytest.mark.parametrize("x, mu, var, y", [_make_montecarlo_x_mu_var_y()])
def test_montecarlo_predict_log_density(
    likelihood_var: float, x: TensorType, mu: TensorType, var: TensorType, y: TensorType
) -> None:
    likelihood_gaussian_mc, likelihood_gaussian = _make_montecarlo_likelihoods(likelihood_var)
    assert_allclose(
        likelihood_gaussian_mc.predict_log_density(x, mu, var, y),
        likelihood_gaussian.predict_log_density(x, mu, var, y),
        rtol=5e-4,
        atol=1e-4,
    )


@pytest.mark.parametrize("likelihood_var", [0.3, 0.5, 1.0])
@pytest.mark.parametrize("x, mu, var, y", [_make_montecarlo_x_mu_var_y()])
def test_montecarlo_predict_mean_and_var(
    likelihood_var: float, x: TensorType, mu: TensorType, var: TensorType, y: TensorType
) -> None:
    likelihood_gaussian_mc, likelihood_gaussian = _make_montecarlo_likelihoods(likelihood_var)
    mean1, var1 = likelihood_gaussian_mc.predict_mean_and_var(x, mu, var)
    mean2, var2 = likelihood_gaussian.predict_mean_and_var(x, mu, var)
    assert_allclose(mean1, mean2, rtol=5e-4, atol=1e-4)
    assert_allclose(var1, var2, rtol=5e-4, atol=1e-4)


# tests/gpflow/likelihoods/test_multiclass.py

# Copyright 2017-2020 the GPflow authors.
#
# 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.

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

from gpflow.base import AnyNDArray, TensorType
from gpflow.config import default_float, default_int
from gpflow.experimental.check_shapes import inherit_check_shapes
from gpflow.likelihoods import Bernoulli, MultiClass, RobustMax, Softmax
from gpflow.utilities import to_default_float, to_default_int

tf.random.set_seed(99012)


@pytest.mark.parametrize("dimX", [3])
@pytest.mark.parametrize("num, dimF", [[10, 5], [3, 2]])
@pytest.mark.parametrize("dimY", [10, 2, 1])
def test_softmax_y_shape_assert(num: int, dimX: int, dimF: int, dimY: int) -> None:
    """
    SoftMax assumes the class is given as a label (not, e.g., one-hot
    encoded), and hence just uses the first column of Y. To prevent
    silent errors, there is a tf assertion that ensures Y only has one
    dimension. This test checks that this assert works as intended.
    """
    X = tf.random.normal((num, dimX))
    F = tf.random.normal((num, dimF))
    dY: AnyNDArray = np.vstack((np.random.randn(num - 3, dimY), np.ones((3, dimY)))) > 0
    Y = tf.convert_to_tensor(dY, dtype=default_int())
    likelihood = Softmax(dimF)
    try:
        likelihood.log_prob(X, F, Y)
    except tf.errors.InvalidArgumentError as e:
        assert "Condition x == y did not hold." in e.message


@pytest.mark.parametrize("dimX", [3])
@pytest.mark.parametrize("num", [10, 3])
@pytest.mark.parametrize("dimF, dimY", [[2, 1]])
def test_softmax_bernoulli_equivalence(num: int, dimX: int, dimF: int, dimY: int) -> None:
    X = tf.random.normal((num, dimX))
    dF: AnyNDArray = np.vstack(
        (np.random.randn(num - 3, dimF), np.array([[-3.0, 0.0], [3, 0.0], [0.0, 0.0]]))
    )
    dY: AnyNDArray = np.vstack((np.random.randn(num - 3, dimY), np.ones((3, dimY)))) > 0
    F = to_default_float(dF)
    Fvar = tf.exp(tf.stack([F[:, 1], -10.0 + tf.zeros(F.shape[0], dtype=F.dtype)], axis=1))
    F = tf.stack([F[:, 0], tf.zeros(F.shape[0], dtype=F.dtype)], axis=1)
    Y = to_default_int(dY)
    Ylabel = 1 - Y

    softmax_likelihood = Softmax(dimF)
    bernoulli_likelihood = Bernoulli(invlink=tf.sigmoid)
    softmax_likelihood.num_monte_carlo_points = int(
        0.3e7
    )  # Minimum number of points to pass the test on CircleCI
    bernoulli_likelihood.num_gauss_hermite_points = 40

    assert_allclose(
        softmax_likelihood.conditional_mean(X, F)[:, :1],
        bernoulli_likelihood.conditional_mean(X, F[:, :1]),
    )

    assert_allclose(
        softmax_likelihood.conditional_variance(X, F)[:, :1],
        bernoulli_likelihood.conditional_variance(X, F[:, :1]),
    )

    assert_allclose(
        softmax_likelihood.log_prob(X, F, Ylabel),
        bernoulli_likelihood.log_prob(X, F[:, :1], Y.numpy()),
    )

    mean1, var1 = softmax_likelihood.predict_mean_and_var(X, F, Fvar)
    mean2, var2 = bernoulli_likelihood.predict_mean_and_var(X, F[:, :1], Fvar[:, :1])

    assert_allclose(mean1[:, 0, None], mean2, rtol=2e-3)
    assert_allclose(var1[:, 0, None], var2, rtol=2e-3)

    ls_ve = softmax_likelihood.variational_expectations(X, F, Fvar, Ylabel)
    lb_ve = bernoulli_likelihood.variational_expectations(X, F[:, :1], Fvar[:, :1], Y.numpy())
    assert_allclose(ls_ve, lb_ve, rtol=5e-3)


@pytest.mark.parametrize("num_classes, num_points", [[10, 3]])
@pytest.mark.parametrize("tol, epsilon", [[1e-4, 1e-3]])
def test_robust_max_multiclass_symmetric(
    num_classes: int, num_points: int, tol: float, epsilon: float
) -> None:
    """
    This test is based on the observation that for
    symmetric inputs the class predictions must have equal probability.
    """
    rng = np.random.RandomState(1)
    p = 1.0 / num_classes
    X = tf.ones((num_points, 1), dtype=default_float())
    F = tf.ones((num_points, num_classes), dtype=default_float())
    Y = tf.convert_to_tensor(rng.randint(num_classes, size=(num_points, 1)), dtype=default_float())

    likelihood = MultiClass(num_classes)
    likelihood.invlink.epsilon = tf.convert_to_tensor(epsilon, dtype=default_float())

    mu, _ = likelihood.predict_mean_and_var(X, F, F)
    pred = likelihood.predict_log_density(X, F, F, Y)
    variational_expectations = likelihood.variational_expectations(X, F, F, Y)

    expected_mu: AnyNDArray = (
        p * (1.0 - epsilon) + (1.0 - p) * epsilon / (num_classes - 1)
    ) * np.ones((num_points, 1))
    expected_log_density = np.log(expected_mu)

    # assert_allclose() would complain about shape mismatch
    assert np.allclose(mu, expected_mu, tol, tol)
    assert np.allclose(pred, expected_log_density, 1e-3, 1e-3)

    validation_variational_expectation = p * np.log(1.0 - epsilon) + (1.0 - p) * np.log(
        epsilon / (num_classes - 1)
    )
    assert_allclose(
        variational_expectations,
        np.ones((num_points,)) * validation_variational_expectation,
        tol,
        tol,
    )


@pytest.mark.parametrize("num_classes, num_points", [[5, 100]])
@pytest.mark.parametrize(
    "mock_prob, expected_prediction, tol, epsilon",
    [
        [0.73, -0.5499780059, 1e-4, 0.231]
        # Expected prediction evaluated on calculator:
        # log((1 - ε) * 0.73 + (1-0.73) * ε / (num_classes -1))
    ],
)
def test_robust_max_multiclass_predict_log_density(
    num_classes: int,
    num_points: int,
    mock_prob: float,
    expected_prediction: float,
    tol: float,
    epsilon: float,
) -> None:
    class MockRobustMax(RobustMax):
        @inherit_check_shapes
        def prob_is_largest(
            self,
            Y: TensorType,
            mu: TensorType,
            var: TensorType,
            gh_x: TensorType,
            gh_w: TensorType,
        ) -> tf.Tensor:
            return tf.ones((num_points, 1), dtype=default_float()) * mock_prob

    likelihood = MultiClass(num_classes, invlink=MockRobustMax(num_classes, epsilon))
    X = tf.ones((num_points, 2))
    F = tf.ones((num_points, num_classes))
    rng = np.random.RandomState(1)
    Y = to_default_int(rng.randint(num_classes, size=(num_points, 1)))
    prediction = likelihood.predict_log_density(X, F, F, Y)

    assert_allclose(prediction, expected_prediction, tol, tol)


@pytest.mark.parametrize("num_classes", [5, 100])
@pytest.mark.parametrize("initial_epsilon, new_epsilon", [[1e-3, 0.412]])
def test_robust_max_multiclass_eps_k1_changes(
    num_classes: int, initial_epsilon: float, new_epsilon: float
) -> None:
    """
    Checks that eps K1 changes when epsilon changes. This used to not happen and had to be
    manually changed.
    """
    likelihood = RobustMax(num_classes, initial_epsilon)
    expected_eps_k1 = initial_epsilon / (num_classes - 1.0)
    actual_eps_k1 = likelihood.eps_k1
    assert_allclose(expected_eps_k1, actual_eps_k1)

    likelihood.epsilon = tf.convert_to_tensor(new_epsilon, dtype=default_float())
    expected_eps_k2 = new_epsilon / (num_classes - 1.0)
    actual_eps_k2 = likelihood.eps_k1
    assert_allclose(expected_eps_k2, actual_eps_k2)


@pytest.mark.skip(
    "ndiagquad cannot handle MultiClass (see https://github.com/GPflow/GPflow/issues/1091"
)
def test_multiclass_quadrature_variational_expectations() -> None:
    pass


@pytest.mark.skip(
    "ndiagquad cannot handle MultiClass (see https://github.com/GPflow/GPflow/issues/1091"
)
def test_multiclass_quadrature_predict_log_density() -> None:
    pass


@pytest.mark.skip(
    "ndiagquad cannot handle MultiClass (see https://github.com/GPflow/GPflow/issues/1091"
)
def test_multiclass_quadrature_predict_mean_and_var() -> None:
    pass


# tests/gpflow/likelihoods/test_switched_likelihood.py

# Copyright 2020 the GPflow authors.
#
# 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 typing import Sequence

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
from gpflow.base import AnyNDArray, Parameter, TensorType
from gpflow.inducing_variables import InducingPoints
from gpflow.likelihoods import Gaussian, StudentT, SwitchedLikelihood


@pytest.mark.parametrize("X_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("Y_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("F_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("Fvar_list", [[tf.exp(tf.random.normal((i, 2))) for i in range(3, 6)]])
@pytest.mark.parametrize("Y_label", [[tf.ones((i, 2)) * (i - 3.0) for i in range(3, 6)]])
def test_switched_likelihood_log_prob(
    X_list: Sequence[TensorType],
    Y_list: Sequence[TensorType],
    F_list: Sequence[TensorType],
    Fvar_list: Sequence[TensorType],
    Y_label: Sequence[TensorType],
) -> None:
    """
    SwitchedLikelihood is separately tested here.
    Here, we make sure the partition-stitch works fine.
    """
    Y_perm = list(range(3 + 4 + 5))
    np.random.shuffle(Y_perm)
    # shuffle the original data
    X_sw = np.concatenate(X_list)[Y_perm, :]
    Y_sw = np.hstack([np.concatenate(Y_list), np.concatenate(Y_label)])[Y_perm, :3]
    F_sw = np.concatenate(F_list)[Y_perm, :]
    likelihoods = [Gaussian()] * 3
    for lik in likelihoods:
        lik.variance = Parameter(np.exp(np.random.randn()), dtype=np.float32)
    switched_likelihood = SwitchedLikelihood(likelihoods)

    switched_results = switched_likelihood.log_prob(X_sw, F_sw, Y_sw)
    results = [lik.log_prob(x, f, y) for lik, x, y, f in zip(likelihoods, X_list, Y_list, F_list)]

    assert_allclose(switched_results, np.concatenate(results)[Y_perm])


@pytest.mark.parametrize("X_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("Y_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("F_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("Fvar_list", [[tf.exp(tf.random.normal((i, 2))) for i in range(3, 6)]])
@pytest.mark.parametrize("Y_label", [[tf.ones((i, 2)) * (i - 3.0) for i in range(3, 6)]])
def test_switched_likelihood_predict_log_density(
    X_list: Sequence[TensorType],
    Y_list: Sequence[TensorType],
    F_list: Sequence[TensorType],
    Fvar_list: Sequence[TensorType],
    Y_label: Sequence[TensorType],
) -> None:
    Y_perm = list(range(3 + 4 + 5))
    np.random.shuffle(Y_perm)
    # shuffle the original data
    X_sw = np.concatenate(X_list)[Y_perm, :]
    Y_sw = np.hstack([np.concatenate(Y_list), np.concatenate(Y_label)])[Y_perm, :3]
    F_sw = np.concatenate(F_list)[Y_perm, :]
    Fvar_sw = np.concatenate(Fvar_list)[Y_perm, :]

    likelihoods = [Gaussian()] * 3
    for lik in likelihoods:
        lik.variance = Parameter(np.exp(np.random.randn()), dtype=np.float32)
    switched_likelihood = SwitchedLikelihood(likelihoods)

    switched_results = switched_likelihood.predict_log_density(X_sw, F_sw, Fvar_sw, Y_sw)
    # likelihood
    results = [
        lik.predict_log_density(x, f, fvar, y)
        for lik, x, y, f, fvar in zip(likelihoods, X_list, Y_list, F_list, Fvar_list)
    ]
    assert_allclose(switched_results, np.concatenate(results)[Y_perm])


@pytest.mark.parametrize("X_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("Y_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("F_list", [[tf.random.normal((i, 2)) for i in range(3, 6)]])
@pytest.mark.parametrize("Fvar_list", [[tf.exp(tf.random.normal((i, 2))) for i in range(3, 6)]])
@pytest.mark.parametrize("Y_label", [[tf.ones((i, 2)) * (i - 3.0) for i in range(3, 6)]])
def test_switched_likelihood_variational_expectations(
    X_list: Sequence[TensorType],
    Y_list: Sequence[TensorType],
    F_list: Sequence[TensorType],
    Fvar_list: Sequence[TensorType],
    Y_label: Sequence[TensorType],
) -> None:
    Y_perm = list(range(3 + 4 + 5))
    np.random.shuffle(Y_perm)
    # shuffle the original data
    X_sw = np.concatenate(X_list)[Y_perm, :]
    Y_sw = np.hstack([np.concatenate(Y_list), np.concatenate(Y_label)])[Y_perm, :3]
    F_sw = np.concatenate(F_list)[Y_perm, :]
    Fvar_sw = np.concatenate(Fvar_list)[Y_perm, :]

    likelihoods = [Gaussian()] * 3
    for lik in likelihoods:
        lik.variance = Parameter(np.exp(np.random.randn()), dtype=np.float32)
    switched_likelihood = SwitchedLikelihood(likelihoods)

    switched_results = switched_likelihood.variational_expectations(X_sw, F_sw, Fvar_sw, Y_sw)
    results = [
        lik.variational_expectations(x, f, fvar, y)
        for lik, x, y, f, fvar in zip(likelihoods, X_list, Y_list, F_list, Fvar_list)
    ]
    assert_allclose(switched_results, np.concatenate(results)[Y_perm])


def test_switched_likelihood_with_vgp() -> None:
    """
    Reproduces the bug in https://github.com/GPflow/GPflow/issues/951
    """
    X = np.random.randn(12 + 15, 1)
    Y = np.random.randn(12 + 15, 1)
    idx: AnyNDArray = np.array([0] * 12 + [1] * 15)
    Y_aug = np.c_[Y, idx]
    assert Y_aug.shape == (12 + 15, 2)

    kernel = gpflow.kernels.Matern32()
    likelihood = gpflow.likelihoods.SwitchedLikelihood([StudentT(), StudentT()])
    model = gpflow.models.VGP((X, Y_aug), kernel=kernel, likelihood=likelihood)
    # without bugfix, optimization errors out
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables, options=dict(maxiter=1))


@pytest.mark.parametrize("num_latent_gps", [1])
def test_switched_likelihood_regression_valid_num_latent_gps(num_latent_gps: int) -> None:
    """
    A Regression test when using Switched likelihood: the number of latent
    functions in a GP model must be equal to the number of columns in Y minus
    one. The final column of Y is used to index the switch. If the number of
    latent functions does not match, an exception will be raised.
    """
    x = np.random.randn(100, 1)
    y: AnyNDArray = np.hstack((np.random.randn(100, 1), np.random.randint(0, 3, (100, 1))))
    data = x, y

    Z = InducingPoints(np.random.randn(num_latent_gps, 1))
    likelihoods = [StudentT()] * 3
    switched_likelihood = SwitchedLikelihood(likelihoods)
    m = gpflow.models.SVGP(
        kernel=gpflow.kernels.Matern12(),
        inducing_variable=Z,
        likelihood=switched_likelihood,
        num_latent_gps=num_latent_gps,
    )
    m.training_loss(data)


# tests/gpflow/models/test_cglb.py

# Copyright 2021 the GPflow authors.
#
# 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 typing import Tuple, cast

import numpy as np
import tensorflow as tf

from gpflow.base import AnyNDArray, RegressionData
from gpflow.config import default_float
from gpflow.experimental.check_shapes import check_shapes
from gpflow.kernels import SquaredExponential
from gpflow.models import CGLB, GPR, SGPR
from gpflow.models.cglb import NystromPreconditioner, cglb_conjugate_gradient
from gpflow.utilities import to_default_float as tdf


@check_shapes(
    "return[0][0]: [N, D]  # X",
    "return[0][1]: [N, P]  # Y",
    "return[1]: [M, D]  # Z, inducing points",
    "return[2]: [Nnew, D]  # Xnew",
)
def data(rng: np.random.RandomState) -> Tuple[RegressionData, tf.Tensor, tf.Tensor]:
    n: int = 100
    t: int = 20
    d: int = 2

    x: AnyNDArray = rng.randn(n, d)
    xs: AnyNDArray = rng.randn(t, d)  # test points
    c: AnyNDArray = np.array([[-1.4], [0.5]])
    y = np.sin(cast(AnyNDArray, x @ c) + 0.5 * rng.randn(n, 1))
    z = rng.randn(10, 2)

    return (tdf(x), tdf(y)), tdf(z), tdf(xs)


def test_cglb_check_basics() -> None:
    """
    * Quadratic term of CGLB with v=0 is equivalent to the quadratic term of SGPR.
    * Log determinant term of CGLB is less or equal to SGPR log determinant.
        In the test the `logdet_term` method returns negative half of the logdet bound,
        therefore we run the opposite direction of the sign.
    """

    rng: np.random.RandomState = np.random.RandomState(999)
    train, z, _ = data(rng)
    noise = 0.2

    sgpr = SGPR(train, kernel=SquaredExponential(), inducing_variable=z, noise_variance=noise)

    # `v_grad_optimization=True` turns off the CG in the quadratic term
    cglb = CGLB(
        train,
        kernel=SquaredExponential(),
        inducing_variable=z,
        noise_variance=noise,
        v_grad_optimization=True,
    )

    sgpr_common = sgpr._common_calculation()
    cglb_common = cglb._common_calculation()

    sgpr_quad_term = sgpr.quad_term(sgpr_common)
    cglb_quad_term = cglb.quad_term(cglb_common)
    np.testing.assert_almost_equal(sgpr_quad_term, cglb_quad_term)

    sgpr_logdet = sgpr.logdet_term(sgpr_common)
    cglb_logdet = cglb.logdet_term(cglb_common)
    assert cglb_logdet >= sgpr_logdet

    x = train[0]
    K = SquaredExponential()(x) + noise * tf.eye(x.shape[0], dtype=default_float())
    gpr_logdet = -0.5 * tf.linalg.logdet(K)
    assert cglb_logdet <= gpr_logdet


def test_conjugate_gradient_convergence() -> None:
    """
    Check that the method of conjugate gradients implemented can solve a linear system of equations
    """
    rng: np.random.RandomState = np.random.RandomState(999)
    noise = 1e-3
    train, z, _ = data(rng)
    x, y = train
    n = x.shape[0]
    b = tf.transpose(y)
    k = SquaredExponential()
    K = k(x) + noise * tf.eye(n, dtype=default_float())
    Kinv_y = tf.linalg.solve(K, y)  # We could solve by cholesky instead

    model = CGLB((x, y), kernel=k, inducing_variable=z, noise_variance=noise)
    common = model._common_calculation()

    initial = tf.zeros_like(b)
    A = common.A
    LB = common.LB
    max_error = 0.01
    max_steps = 200
    restart_cg_step = 200
    preconditioner = NystromPreconditioner(A, LB, noise)

    v = cglb_conjugate_gradient(
        K, b, initial, preconditioner, max_error, max_steps, restart_cg_step
    )

    # NOTE: with smaller `max_error` we can reduce the `rtol`
    np.testing.assert_allclose(Kinv_y, tf.transpose(v), rtol=0.1)


def test_cglb_quad_term_guarantees() -> None:
    """
    Check that when conjugate gradient is used to evaluate the quadratic term,
    the obtained solution is:

    1. Smaller than the solution computed by Cholesky decomposition
    2. Within the error tolerance of the solution computed by Cholesky
    """
    rng: np.random.RandomState = np.random.RandomState(999)

    max_error: float = 1e-2
    noise: float = 1e-2
    train, z, _ = data(rng)
    x, y = train
    k = SquaredExponential()
    K = k(x) + noise * tf.eye(x.shape[0], dtype=default_float())

    def inv_quad_term(K: tf.Tensor, y: tf.Tensor) -> tf.Tensor:
        """
        For PSD K, compute -0.5 * y.T K^{-1} y via Cholesky decomposition
        """
        L = tf.linalg.cholesky(K)
        Linvy = tf.linalg.triangular_solve(L, y)
        return -0.5 * tf.reduce_sum(tf.square(Linvy))

    cholesky_quad_term = inv_quad_term(K, y)

    cglb = CGLB(
        train,
        kernel=k,
        inducing_variable=z,
        noise_variance=noise,
        cg_tolerance=max_error,
        max_cg_iters=100,
        restart_cg_iters=10,
    )

    common = cglb._common_calculation()
    cglb_quad_term = cglb.quad_term(common)

    assert cglb_quad_term <= cholesky_quad_term
    assert np.abs(cglb_quad_term - cholesky_quad_term) <= max_error


def test_cglb_predict() -> None:
    """
    Test that 1.) The predict method returns the same variance estimate as SGPR.
              2.) The predict method returns the same mean as SGPR for v=0.
              3.) The predict method returns a mean very similar to GPR when CG is run to low tolerance.
    """
    rng: np.random.RandomState = np.random.RandomState(999)
    train, z, xs = data(rng)
    noise = 0.2

    gpr = GPR(train, kernel=SquaredExponential(), noise_variance=noise)
    sgpr = SGPR(train, kernel=SquaredExponential(), inducing_variable=z, noise_variance=noise)

    cglb = CGLB(
        train,
        kernel=SquaredExponential(),
        inducing_variable=z,
        noise_variance=noise,
    )

    gpr_mean, _ = gpr.predict_y(xs, full_cov=False)
    sgpr_mean, sgpr_cov = sgpr.predict_y(xs, full_cov=False)
    cglb_mean, cglb_cov = cglb.predict_y(
        xs, full_cov=False, cg_tolerance=1e6
    )  # set tolerance high so v stays at 0.

    assert np.allclose(sgpr_cov, cglb_cov)
    assert np.allclose(sgpr_mean, cglb_mean)

    cglb_mean, _ = cglb.predict_y(xs, full_cov=False, cg_tolerance=1e-12)

    assert np.allclose(gpr_mean, cglb_mean)


# tests/gpflow/models/test_gplvm.py

# Copyright 2019 the GPflow authors.
#
# 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 typing import Optional

import numpy as np
import pytest

import gpflow
from gpflow.utilities.ops import pca_reduce


class Data:
    rng = np.random.RandomState(999)
    N = 20
    D = 5
    Y = rng.randn(N, D)
    Q = 2
    M = 10
    X = rng.randn(N, Q)


@pytest.mark.parametrize(
    "kernel",
    [
        None,  # default kernel: SquaredExponential
        gpflow.kernels.Periodic(base_kernel=gpflow.kernels.SquaredExponential()),
    ],
)
def test_gplvm_with_kernels(kernel: Optional[gpflow.kernels.Kernel]) -> None:
    m = gpflow.models.GPLVM(Data.Y, Data.Q, kernel=kernel)
    lml_initial = m.log_marginal_likelihood()
    opt = gpflow.optimizers.Scipy()
    opt.minimize(m.training_loss, m.trainable_variables, options=dict(maxiter=2))
    assert m.log_marginal_likelihood() > lml_initial


def test_bayesian_gplvm_1d() -> None:
    Q = 1
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = np.linspace(0, 1, Data.M)[:, None]
    m = gpflow.models.BayesianGPLVM(
        Data.Y,
        np.zeros((Data.N, Q)),
        np.ones((Data.N, Q)),
        kernel,
        inducing_variable=inducing_variable,
    )
    assert m.inducing_variable.num_inducing == Data.M

    elbo_initial = m.elbo()
    opt = gpflow.optimizers.Scipy()
    opt.minimize(m.training_loss, m.trainable_variables, options=dict(maxiter=2))
    assert m.elbo() > elbo_initial


def test_bayesian_gplvm_2d() -> None:
    Q = 2  # latent dimensions
    X_data_mean = pca_reduce(Data.Y, Q)
    kernel = gpflow.kernels.SquaredExponential()

    m = gpflow.models.BayesianGPLVM(
        Data.Y, X_data_mean, np.ones((Data.N, Q)), kernel, num_inducing_variables=Data.M
    )

    elbo_initial = m.elbo()
    opt = gpflow.optimizers.Scipy()
    opt.minimize(m.training_loss, m.trainable_variables, options=dict(maxiter=2))
    assert m.elbo() > elbo_initial

    # test prediction
    Xtest = Data.rng.randn(10, Q)
    mu_f, var_f = m.predict_f(Xtest)
    mu_fFull, var_fFull = m.predict_f(Xtest, full_cov=True)
    np.testing.assert_allclose(mu_fFull, mu_f)

    for i in range(Data.D):
        np.testing.assert_allclose(var_f[:, i], np.diag(var_fFull[i, :, :]))


def test_gplvm_constructor_checks() -> None:
    with pytest.raises(ValueError):
        assert Data.X.shape[1] == Data.Q
        latents_wrong_shape = Data.X[:, : Data.Q - 1]
        gpflow.models.GPLVM(Data.Y, Data.Q, X_data_mean=latents_wrong_shape)
    with pytest.raises(ValueError):
        observations_wrong_shape = Data.Y[:, : Data.Q - 1]
        gpflow.models.GPLVM(observations_wrong_shape, Data.Q)
    with pytest.raises(ValueError):
        observations_wrong_shape = Data.Y[:, : Data.Q - 1]
        gpflow.models.GPLVM(observations_wrong_shape, Data.Q, X_data_mean=Data.X)


def test_bayesian_gplvm_constructor_check() -> None:
    Q = 1
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = np.linspace(0, 1, Data.M)[:, None]
    with pytest.raises(ValueError):
        gpflow.models.BayesianGPLVM(
            Data.Y,
            np.zeros((Data.N, Q)),
            np.ones((Data.N, Q)),
            kernel,
            inducing_variable=inducing_variable,
            num_inducing_variables=len(inducing_variable),
        )


# tests/gpflow/models/test_gpr.py

# Copyright 2019 the GPflow authors.
#
# 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.

import numpy as np
import tensorflow as tf

import gpflow
from gpflow import set_trainable

rng = np.random.RandomState(0)


class Data:
    N = 10
    D = 1
    X = rng.rand(N, D)
    Y = rng.rand(N, 1)
    ls = 2.0
    var = 1.0


def test_non_trainable_model_objective() -> None:
    """
    Checks that we can still compute the objective of a model that has no
    trainable parameters whatsoever (regression test for bug in log_prior()).
    In this case we have no priors, so log_prior should be zero to add no
    contribution to the objective.
    """
    model = gpflow.models.GPR(
        (Data.X, Data.Y),
        kernel=gpflow.kernels.SquaredExponential(lengthscales=Data.ls, variance=Data.var),
    )

    set_trainable(model, False)

    _ = model.log_marginal_likelihood()
    assert model.log_prior_density() == 0.0


def test_varying_data() -> None:
    input_dim = 2
    output_dim = 1

    N = 5
    X, Y = rng.randn(N, input_dim), rng.randn(N, output_dim)

    var_data = (tf.Variable(X, shape=[None, input_dim]), tf.Variable(Y, shape=[None, output_dim]))
    m = gpflow.models.GPR(var_data, gpflow.kernels.SquaredExponential())

    lml_func = tf.function(m.log_marginal_likelihood)
    old_lml = lml_func()

    new_N = 7
    new_X, new_Y = rng.randn(new_N, input_dim), rng.randn(new_N, output_dim)

    # assign new data:
    for var, new_value in zip(var_data, (new_X, new_Y)):
        var.assign(new_value)

    new_lml = lml_func()  # re-use compiled function

    assert (
        np.abs((old_lml - new_lml) / (new_lml + old_lml)) > 0.1
    ), "we expect the LML for different data to be significantly different"


# tests/gpflow/models/test_gpr_posterior.py

from typing import Tuple

import numpy as np
import pytest

import gpflow
from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import check_shapes
from gpflow.models.gpr import GPR_deprecated, GPR_with_posterior
from gpflow.posteriors import PrecomputeCacheType


@check_shapes(
    "regression_data[0]: [N, D]",
    "regression_data[1]: [N, P]",
)
def make_models(
    regression_data: gpflow.base.RegressionData,
) -> Tuple[GPR_deprecated, GPR_with_posterior]:
    """Helper function to create models"""

    k = gpflow.kernels.Matern52()

    mold = GPR_deprecated(data=regression_data, kernel=k)
    mnew = GPR_with_posterior(data=regression_data, kernel=k)
    return mold, mnew


@check_shapes(
    "return[0]: [N, D]  # X",
    "return[1]: [batch_new..., N_new, D]  # X_new",
    "return[2]: [N, P]  # Y",
)
def _get_data_for_tests() -> Tuple[AnyNDArray, AnyNDArray, AnyNDArray]:
    """Helper function to create testing data"""
    X = np.random.randn(5, 6)
    Y = np.random.randn(5, 2)
    X_new = np.random.randn(3, 10, 5, 6)
    return X, X_new, Y


@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [False])
def test_old_vs_new_gp_fused(full_cov: bool, full_output_cov: bool) -> None:
    X, X_new, Y = _get_data_for_tests()
    mold, mnew = make_models((X, Y))

    mu_old, var2_old = mold.predict_f(X_new, full_cov=full_cov, full_output_cov=full_output_cov)
    mu_new_fuse, var2_new_fuse = mnew.predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )
    # check new fuse is same as old version
    np.testing.assert_allclose(mu_new_fuse, mu_old)
    np.testing.assert_allclose(var2_new_fuse, var2_old)


@pytest.mark.parametrize("cache_type", [PrecomputeCacheType.TENSOR, PrecomputeCacheType.VARIABLE])
@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [False])
def test_old_vs_new_with_posterior(
    cache_type: PrecomputeCacheType, full_cov: bool, full_output_cov: bool
) -> None:
    X, X_new, Y = _get_data_for_tests()
    mold, mnew = make_models((X, Y))

    mu_old, var2_old = mold.predict_f(X_new, full_cov=full_cov, full_output_cov=full_output_cov)
    mu_new_cache, var2_new_cache = mnew.posterior(cache_type).predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )

    # check new cache is same as old version
    np.testing.assert_allclose(mu_old, mu_new_cache)
    np.testing.assert_allclose(var2_old, var2_new_cache)


# tests/gpflow/models/test_mcmc.py

# Copyright 2016-2020 the GPflow authors.
#
# 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.

import numpy as np
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
from gpflow.config import default_float


def test_sparse_mcmc_likelihoods_and_gradients() -> None:
    """
    This test makes sure that when the inducing points are the same as the data
    points, the sparse mcmc is the same as full mcmc
    """
    rng = np.random.RandomState(0)
    X, Y = rng.randn(10, 1), rng.randn(10, 1)
    v_vals = rng.randn(10, 1)

    likelihood = gpflow.likelihoods.StudentT()
    model_1 = gpflow.models.GPMC(
        data=(X, Y), kernel=gpflow.kernels.Exponential(), likelihood=likelihood
    )
    model_2 = gpflow.models.SGPMC(
        data=(X, Y),
        kernel=gpflow.kernels.Exponential(),
        inducing_variable=X.copy(),
        likelihood=likelihood,
    )
    model_1.V = tf.convert_to_tensor(v_vals, dtype=default_float())
    model_2.V = tf.convert_to_tensor(v_vals, dtype=default_float())
    model_1.kernel.lengthscales.assign(0.8)
    model_2.kernel.lengthscales.assign(0.8)
    model_1.kernel.variance.assign(4.2)
    model_2.kernel.variance.assign(4.2)

    assert_allclose(
        model_1.log_posterior_density(), model_2.log_posterior_density(), rtol=1e-5, atol=1e-5
    )


# tests/gpflow/models/test_methods.py

# Copyright 2016 the GPflow authors
#
# 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 dataclasses import dataclass
from typing import List

import numpy as np
import pytest
from numpy.testing import assert_array_equal, assert_array_less

import gpflow
from gpflow.base import AnyNDArray

# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


@dataclass(frozen=True)
class Datum:
    rng: np.random.RandomState = np.random.RandomState(0)
    X: AnyNDArray = rng.randn(100, 2)
    Y: AnyNDArray = rng.randn(100, 1)
    Z: AnyNDArray = rng.randn(10, 2)
    Xs: AnyNDArray = rng.randn(10, 2)
    lik = gpflow.likelihoods.Gaussian()
    kernel = gpflow.kernels.Matern32()


default_datum = Datum()


_gp_models: List[gpflow.models.GPModel] = [
    gpflow.models.VGP((default_datum.X, default_datum.Y), default_datum.kernel, default_datum.lik),
    gpflow.models.GPMC((default_datum.X, default_datum.Y), default_datum.kernel, default_datum.lik),
    gpflow.models.SGPMC(
        (default_datum.X, default_datum.Y),
        default_datum.kernel,
        default_datum.lik,
        inducing_variable=default_datum.Z,
    ),
    gpflow.models.SGPR(
        (default_datum.X, default_datum.Y),
        default_datum.kernel,
        inducing_variable=default_datum.Z,
    ),
    gpflow.models.GPR((default_datum.X, default_datum.Y), default_datum.kernel),
    gpflow.models.GPRFITC(
        (default_datum.X, default_datum.Y),
        default_datum.kernel,
        inducing_variable=default_datum.Z,
    ),
]

_state_less_gp_models: List[gpflow.models.GPModel] = [
    gpflow.models.SVGP(default_datum.kernel, default_datum.lik, inducing_variable=default_datum.Z)
]


@pytest.mark.parametrize("model", _state_less_gp_models + _gp_models)
def test_methods_predict_f(model: gpflow.models.GPModel) -> None:
    mf, vf = model.predict_f(default_datum.Xs)
    assert_array_equal(mf.shape, vf.shape)
    assert_array_equal(mf.shape, (10, 1))
    assert_array_less(np.full_like(vf, -1e-6), vf)


@pytest.mark.parametrize("model", _state_less_gp_models + _gp_models)
def test_methods_predict_y(model: gpflow.models.GPModel) -> None:
    mf, vf = model.predict_y(default_datum.Xs)
    assert_array_equal(mf.shape, vf.shape)
    assert_array_equal(mf.shape, (10, 1))
    assert_array_less(np.full_like(vf, -1e-6), vf)


@pytest.mark.parametrize("model", _state_less_gp_models + _gp_models)
def test_methods_predict_log_density(model: gpflow.models.GPModel) -> None:
    rng = Datum().rng
    Ys = rng.randn(10, 1)
    d = model.predict_log_density((default_datum.Xs, Ys))
    assert_array_equal(d.shape, (10,))


# tests/gpflow/models/test_model_predict.py

# Copyright 2016 the GPflow authors.
#
# 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 typing import Any, Dict, Optional, Type

import numpy as np
import pytest

import gpflow
from gpflow.inducing_variables import InducingPoints
from gpflow.kernels import Kernel, Matern32
from gpflow.models.util import InducingPointsLike

rng = np.random.RandomState(0)


class ModelSetup:
    def __init__(
        self,
        model_class: Type[gpflow.models.GPModel],
        kernel: Kernel = Matern32(),
        likelihood: gpflow.likelihoods.Likelihood = gpflow.likelihoods.Gaussian(),
        whiten: Optional[bool] = None,
        q_diag: Optional[bool] = None,
        requires_inducing_variables: bool = True,
        requires_data: bool = False,
        requires_likelihood: bool = True,
    ) -> None:

        self.model_class = model_class
        self.kernel = kernel
        self.likelihood = likelihood
        self.whiten = whiten
        self.q_diag = q_diag
        self.requires_inducing_variables = requires_inducing_variables
        self.requires_data = requires_data
        self.requires_likelihood = requires_likelihood

    def get_model(
        self,
        Z: InducingPointsLike,
        num_latent_gps: int,
        data: Optional[gpflow.base.RegressionData] = None,
    ) -> gpflow.models.GPModel:
        params: Dict[str, Any] = dict(kernel=self.kernel, num_latent_gps=num_latent_gps)

        if self.whiten is not None and self.q_diag is not None:
            params.update(inducing_variable=Z, whiten=self.whiten, q_diag=self.q_diag)

        if self.requires_inducing_variables:
            params.update(dict(inducing_variable=Z))

        if self.requires_data:
            params.update(dict(data=data))

        if self.requires_likelihood:
            params.update(dict(likelihood=self.likelihood))

        return self.model_class(**params)

    def __repr__(self) -> str:
        return f"ModelSetup({self.model_class.__name__}, {self.whiten}, {self.q_diag})"


model_setups = [
    ModelSetup(model_class=gpflow.models.SVGP, whiten=False, q_diag=True),
    ModelSetup(model_class=gpflow.models.SVGP, whiten=True, q_diag=False),
    ModelSetup(model_class=gpflow.models.SVGP, whiten=True, q_diag=True),
    ModelSetup(model_class=gpflow.models.SVGP, whiten=False, q_diag=False),
    ModelSetup(
        model_class=gpflow.models.SGPR,  # type: ignore[misc]  # mypy thinks SPGR is abstract.
        requires_data=True,
        requires_likelihood=False,
    ),
    ModelSetup(
        model_class=gpflow.models.VGP,
        requires_inducing_variables=False,
        requires_data=True,
    ),
    #     ModelSetup(model_class=gpflow.models.GPRF),
    ModelSetup(
        model_class=gpflow.models.GPMC,
        requires_data=True,
        requires_inducing_variables=False,
    ),
    ModelSetup(
        model_class=gpflow.models.SGPMC,
        requires_data=True,
        requires_inducing_variables=True,
    ),
]


@pytest.mark.parametrize("Ntrain, Ntest, D", [[100, 10, 2]])
def test_gaussian_mean_and_variance(Ntrain: int, Ntest: int, D: int) -> None:
    data = rng.randn(Ntrain, D), rng.randn(Ntrain, 1)
    Xtest, _ = rng.randn(Ntest, D), rng.randn(Ntest, 1)
    kernel = Matern32() + gpflow.kernels.White()
    model_gp = gpflow.models.GPR(data, kernel=kernel)

    mu_f, var_f = model_gp.predict_f(Xtest)
    mu_y, var_y = model_gp.predict_y(Xtest)

    assert np.allclose(mu_f, mu_y)
    assert np.allclose(var_f, var_y - 1.0)


@pytest.mark.parametrize("Ntrain, Ntest, D", [[100, 10, 2]])
def test_gaussian_log_density(Ntrain: int, Ntest: int, D: int) -> None:
    data = rng.randn(Ntrain, D), rng.randn(Ntrain, 1)
    Xtest, Ytest = rng.randn(Ntest, D), rng.randn(Ntest, 1)
    kernel = Matern32() + gpflow.kernels.White()
    model_gp = gpflow.models.GPR(data, kernel=kernel)

    mu_y, var_y = model_gp.predict_y(Xtest)
    data = Xtest, Ytest
    log_density = model_gp.predict_log_density(data)
    log_density_hand = np.squeeze(
        -0.5 * np.log(2 * np.pi) - 0.5 * np.log(var_y) - 0.5 * np.square(mu_y - Ytest) / var_y,
        axis=-1,
    )

    assert np.allclose(log_density_hand, log_density)


@pytest.mark.parametrize("input_dim, output_dim, N, Ntest", [[3, 2, 20, 30]])
def test_gaussian_full_cov(input_dim: int, output_dim: int, N: int, Ntest: int) -> None:
    covar_shape = (output_dim, Ntest, Ntest)
    X, Y = rng.randn(N, input_dim), rng.randn(N, output_dim)
    Xtest = rng.randn(Ntest, input_dim)
    kernel = Matern32()
    model_gp = gpflow.models.GPR((X, Y), kernel=kernel)

    mu1, var = model_gp.predict_f(Xtest, full_cov=False)
    mu2, covar = model_gp.predict_f(Xtest, full_cov=True)

    assert np.allclose(mu1, mu2, atol=1.0e-10)
    assert covar.shape == covar_shape
    assert var.shape == (Ntest, output_dim)
    for i in range(output_dim):
        assert np.allclose(var[:, i], np.diag(covar[i, :, :]))


@pytest.mark.parametrize("input_dim, output_dim, N, Ntest, M, num_samples", [[3, 2, 20, 30, 5, 5]])
def test_gaussian_full_cov_samples(
    input_dim: int, output_dim: int, N: int, Ntest: int, M: int, num_samples: int
) -> None:
    samples_shape = (num_samples, Ntest, output_dim)
    X, Y, _ = rng.randn(N, input_dim), rng.randn(N, output_dim), rng.randn(M, input_dim)
    Xtest = rng.randn(Ntest, input_dim)
    kernel = Matern32()
    model_gp = gpflow.models.GPR((X, Y), kernel=kernel)

    samples = model_gp.predict_f_samples(Xtest, num_samples)
    assert samples.shape == samples_shape

    samples = model_gp.predict_f_samples(Xtest, num_samples, full_cov=False)
    assert samples.shape == samples_shape


@pytest.mark.parametrize("model_setup", model_setups)
@pytest.mark.parametrize("input_dim", [3])
@pytest.mark.parametrize("output_dim", [2])
@pytest.mark.parametrize("N", [20])
@pytest.mark.parametrize("Ntest", [30])
@pytest.mark.parametrize("M", [5])
def test_other_models_full_cov(
    model_setup: ModelSetup, input_dim: int, output_dim: int, N: int, Ntest: int, M: int
) -> None:
    covar_shape = (output_dim, Ntest, Ntest)
    X, Y = rng.randn(N, input_dim), rng.randn(N, output_dim)
    Z = InducingPoints(rng.randn(M, input_dim))
    Xtest = rng.randn(Ntest, input_dim)
    model_gp = model_setup.get_model(Z, num_latent_gps=output_dim, data=(X, Y))

    mu1, var = model_gp.predict_f(Xtest, full_cov=False)
    mu2, covar = model_gp.predict_f(Xtest, full_cov=True)

    assert np.allclose(mu1, mu2, atol=1.0e-10)
    assert covar.shape == covar_shape
    assert var.shape == (Ntest, output_dim)
    for i in range(output_dim):
        assert np.allclose(var[:, i], np.diag(covar[i, :, :]))


@pytest.mark.parametrize("model_setup", model_setups)
@pytest.mark.parametrize("input_dim", [3])
@pytest.mark.parametrize("output_dim", [2])
@pytest.mark.parametrize("N", [20])
@pytest.mark.parametrize("Ntest", [30])
@pytest.mark.parametrize("M", [5])
@pytest.mark.parametrize("num_samples", [5])
def test_other_models_full_cov_samples(
    model_setup: ModelSetup,
    input_dim: int,
    output_dim: int,
    N: int,
    Ntest: int,
    M: int,
    num_samples: int,
) -> None:
    samples_shape = (num_samples, Ntest, output_dim)
    X, Y, Z = rng.randn(N, input_dim), rng.randn(N, output_dim), rng.randn(M, input_dim)
    Xtest = rng.randn(Ntest, input_dim)
    model_gp = model_setup.get_model(Z, num_latent_gps=output_dim, data=(X, Y))

    samples = model_gp.predict_f_samples(Xtest, num_samples)
    assert samples.shape == samples_shape


# tests/gpflow/models/test_sgpr.py

# Copyright 2016-2020 the GPflow authors.
#
# 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 dataclasses import dataclass

import numpy as np
import tensorflow as tf

import gpflow
from gpflow.base import AnyNDArray
from gpflow.utilities import to_default_float


@dataclass(frozen=True)
class Datum:
    rng: np.random.RandomState = np.random.RandomState(0)
    X: AnyNDArray = rng.randn(100, 2)
    Y: AnyNDArray = rng.randn(100, 1)
    Z: AnyNDArray = rng.randn(20, 2)


def test_sgpr_qu() -> None:
    rng = np.random.RandomState(1)
    X = to_default_float(Datum.X)
    Z = to_default_float(Datum.Z)
    Y = to_default_float(np.sin(X @ np.array([[-1.4], [0.5]])) + 0.5 * rng.randn(len(X), 1))
    model = gpflow.models.SGPR(
        (X, Y), kernel=gpflow.kernels.SquaredExponential(), inducing_variable=Z
    )

    gpflow.optimizers.Scipy().minimize(model.training_loss, variables=model.trainable_variables)

    qu_mean, qu_cov = model.compute_qu()
    f_at_Z_mean, f_at_Z_cov = model.predict_f(model.inducing_variable.Z, full_cov=True)

    np.testing.assert_allclose(qu_mean, f_at_Z_mean, rtol=1e-5, atol=1e-5)
    np.testing.assert_allclose(tf.reshape(qu_cov, (1, 20, 20)), f_at_Z_cov, rtol=1e-5, atol=1e-5)


def test_sgpr_svgp_qu_equivivalent() -> None:
    rng = np.random.RandomState(2)
    X = to_default_float(Datum.X)
    Y = to_default_float(Datum.Y)
    Z = to_default_float(Datum.Z)
    kernel = gpflow.kernels.SquaredExponential()
    linear = gpflow.functions.Linear(A=[[0.1], [0.2]], b=0.3)
    likelihood = gpflow.likelihoods.Gaussian(scale=linear)
    sgpr = gpflow.models.SGPR(
        (X, Y),
        kernel=kernel,
        inducing_variable=Z,
        likelihood=likelihood,
    )

    q_mu, q_cov = sgpr.compute_qu()
    q_sqrt = tf.linalg.cholesky(q_cov)
    q_sqrt = q_sqrt[None, :, :]
    svgp = gpflow.models.SVGP(
        kernel=kernel,
        likelihood=likelihood,
        inducing_variable=Z,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=False,
    )

    Xnew = to_default_float(rng.randn(100, 2))
    np.testing.assert_allclose(
        sgpr.predict_f(Xnew),
        svgp.predict_f(Xnew),
        atol=1e-4,
    )


# tests/gpflow/models/test_sgpr_posterior.py

# Copyright 2016-2020 the GPflow authors.
#
# 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.
""" Tests the SGPR posterior. """

from typing import Tuple

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import InputData, OutputData
from gpflow.experimental.check_shapes import check_shapes
from gpflow.inducing_variables import InducingPoints
from gpflow.models.sgpr import SGPR, SGPR_deprecated
from gpflow.posteriors import PrecomputeCacheType

INPUT_DIM = 7
OUTPUT_DIM = 1

MEAN_FUNCTION = gpflow.mean_functions.Constant(c=1.0)
KERNEL = gpflow.kernels.Matern52()
Z = np.random.randn(20, INPUT_DIM)


@pytest.fixture(name="dummy_data", scope="module")
@check_shapes(
    "return[0]: [N, D]  # X",
    "return[1]: [N_new, D]  # X_new",
    "return[2]: [N, P]  # Y",
)
def _dummy_data() -> Tuple[InputData, InputData, OutputData]:
    X = tf.convert_to_tensor(np.random.randn(100, INPUT_DIM), dtype=tf.float64)
    Y = tf.convert_to_tensor(np.random.randn(100, OUTPUT_DIM), dtype=tf.float64)
    X_new = tf.convert_to_tensor(np.random.randn(100, INPUT_DIM), dtype=tf.float64)

    return X, X_new, Y


@pytest.fixture(name="sgpr_deprecated_model", scope="module")
@check_shapes(
    "dummy_data[0]: [N, D]  # X",
    "dummy_data[1]: [N_new, D]  # X_new",
    "dummy_data[2]: [N, P]  # Y",
)
def _sgpr_deprecated_model(dummy_data: Tuple[InputData, InputData, OutputData]) -> SGPR_deprecated:
    X, _, Y = dummy_data
    return SGPR_deprecated(
        data=(X, Y), kernel=KERNEL, inducing_variable=InducingPoints(Z), mean_function=MEAN_FUNCTION
    )


@pytest.fixture(name="sgpr_model", scope="module")
@check_shapes(
    "dummy_data[0]: [N, D]  # X",
    "dummy_data[1]: [N_new, D]  # X_new",
    "dummy_data[2]: [N, P]  # Y",
)
def sgpr_model(dummy_data: Tuple[InputData, InputData, OutputData]) -> SGPR:
    X, _, Y = dummy_data
    return SGPR(
        data=(X, Y), kernel=KERNEL, inducing_variable=InducingPoints(Z), mean_function=MEAN_FUNCTION
    )


@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [False])
def test_old_vs_new_gp_fused(
    sgpr_deprecated_model: SGPR_deprecated,
    sgpr_model: SGPR,
    dummy_data: Tuple[InputData, InputData, OutputData],
    full_cov: bool,
    full_output_cov: bool,
) -> None:
    _, X_new, _ = dummy_data

    mu_old, var2_old = sgpr_deprecated_model.predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )
    mu_new_fuse, var2_new_fuse = sgpr_model.predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )

    # check new fuse is same as old version
    np.testing.assert_allclose(mu_new_fuse, mu_old)
    np.testing.assert_allclose(var2_new_fuse, var2_old)


# TODO: move to common test_model_utils
@pytest.mark.parametrize("cache_type", [PrecomputeCacheType.TENSOR, PrecomputeCacheType.VARIABLE])
@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [False])
def test_old_vs_new_with_posterior(
    sgpr_deprecated_model: SGPR_deprecated,
    sgpr_model: SGPR,
    dummy_data: Tuple[InputData, InputData, OutputData],
    cache_type: PrecomputeCacheType,
    full_cov: bool,
    full_output_cov: bool,
) -> None:
    _, X_new, _ = dummy_data

    mu_old, var2_old = sgpr_deprecated_model.predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )
    mu_new_cache, var2_new_cache = sgpr_model.posterior(cache_type).predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )

    # check new cache is same as old version
    np.testing.assert_allclose(mu_old, mu_new_cache)
    np.testing.assert_allclose(var2_old, var2_new_cache)


# tests/gpflow/models/test_svgp.py

# Copyright 2016-2020 the GPflow authors.
#
# 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 dataclasses import dataclass
from typing import Sequence

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow
from gpflow import set_trainable
from gpflow.base import AnyNDArray
from gpflow.models import SVGP


@dataclass(frozen=True)
class DatumSVGP:
    rng: np.random.RandomState = np.random.RandomState(0)
    X: AnyNDArray = rng.randn(20, 1)
    Y: AnyNDArray = rng.randn(20, 2) ** 2
    Z: AnyNDArray = rng.randn(3, 1)
    qsqrt: AnyNDArray = (rng.randn(3, 2) ** 2) * 0.01
    qmean: AnyNDArray = rng.randn(3, 2)
    lik = gpflow.likelihoods.Exponential()
    data = (X, Y)


default_datum_svgp = DatumSVGP()


def test_svgp_fixing_q_sqrt() -> None:
    """
    In response to bug #46, we need to make sure that the q_sqrt matrix can be fixed
    """
    num_latent_gps = default_datum_svgp.Y.shape[1]
    model = SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=default_datum_svgp.lik,
        q_diag=True,
        num_latent_gps=num_latent_gps,
        inducing_variable=default_datum_svgp.Z,
        whiten=False,
    )
    default_num_trainable_variables = len(model.trainable_variables)
    set_trainable(model.q_sqrt, False)
    assert len(model.trainable_variables) == default_num_trainable_variables - 1


def test_svgp_white() -> None:
    """
    Tests that the SVGP bound on the likelihood is the same when using
    with and without diagonals when whitening.
    """
    num_latent_gps = default_datum_svgp.Y.shape[1]
    model_1 = SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=default_datum_svgp.lik,
        q_diag=True,
        num_latent_gps=num_latent_gps,
        inducing_variable=default_datum_svgp.Z,
        whiten=True,
    )
    model_2 = SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=default_datum_svgp.lik,
        q_diag=False,
        num_latent_gps=num_latent_gps,
        inducing_variable=default_datum_svgp.Z,
        whiten=True,
    )
    model_1.q_sqrt.assign(default_datum_svgp.qsqrt)
    model_1.q_mu.assign(default_datum_svgp.qmean)
    model_2.q_sqrt.assign(
        np.array(
            [
                np.diag(default_datum_svgp.qsqrt[:, 0]),
                np.diag(default_datum_svgp.qsqrt[:, 1]),
            ]
        )
    )
    model_2.q_mu.assign(default_datum_svgp.qmean)
    assert_allclose(model_1.elbo(default_datum_svgp.data), model_2.elbo(default_datum_svgp.data))


def test_svgp_non_white() -> None:
    """
    Tests that the SVGP bound on the likelihood is the same when using
    with and without diagonals when whitening is not used.
    """
    num_latent_gps = default_datum_svgp.Y.shape[1]
    model_1 = SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=default_datum_svgp.lik,
        q_diag=True,
        num_latent_gps=num_latent_gps,
        inducing_variable=default_datum_svgp.Z,
        whiten=False,
    )
    model_2 = SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=default_datum_svgp.lik,
        q_diag=False,
        num_latent_gps=num_latent_gps,
        inducing_variable=default_datum_svgp.Z,
        whiten=False,
    )
    model_1.q_sqrt.assign(default_datum_svgp.qsqrt)
    model_1.q_mu.assign(default_datum_svgp.qmean)
    model_2.q_sqrt.assign(
        np.array(
            [
                np.diag(default_datum_svgp.qsqrt[:, 0]),
                np.diag(default_datum_svgp.qsqrt[:, 1]),
            ]
        )
    )
    model_2.q_mu.assign(default_datum_svgp.qmean)
    assert_allclose(model_1.elbo(default_datum_svgp.data), model_2.elbo(default_datum_svgp.data))


def _check_models_close(m1: SVGP, m2: SVGP, tolerance: float = 1e-2) -> bool:
    m1_params = {p.name: p for p in list(m1.trainable_parameters)}
    m2_params = {p.name: p for p in list(m2.trainable_parameters)}
    if set(m2_params.keys()) != set(m2_params.keys()):
        return False
    for key in m1_params:
        p1 = m1_params[key]
        p2 = m2_params[key]
        if not np.allclose(p1.numpy(), p2.numpy(), rtol=tolerance, atol=tolerance):
            return False
    return True


@pytest.mark.parametrize(
    "indices_1, indices_2, num_data1, num_data2, max_iter",
    [
        [[0, 1], [1, 0], 2, 2, 3],
        [[0, 1], [0, 0], 1, 2, 1],
        [[0, 0], [0, 1], 1, 1, 2],
    ],
)
def test_stochastic_gradients(
    indices_1: Sequence[int],
    indices_2: Sequence[int],
    num_data1: int,
    num_data2: int,
    max_iter: int,
) -> None:
    """
    In response to bug #281, we need to make sure stochastic update
    happens correctly in tf optimizer mode.
    To do this compare stochastic updates with deterministic updates
    that should be equivalent.

    Data term in svgp likelihood is
    \sum_{i=1^N}E_{q(i)}[\log p(y_i | f_i )

    This sum is then approximated with an unbiased minibatch estimate.
    In this test we substitute a deterministic analogue of the batchs
    sampler for which we can predict the effects of different updates.
    """
    X: AnyNDArray
    Y: AnyNDArray
    X, Y = np.atleast_2d(np.array([0.0, 1.0])).T, np.atleast_2d(np.array([-1.0, 3.0])).T
    Z: AnyNDArray = np.atleast_2d(np.array([0.5]))

    def get_model(num_data: int) -> SVGP:
        return SVGP(
            kernel=gpflow.kernels.SquaredExponential(),
            num_data=num_data,
            likelihood=gpflow.likelihoods.Gaussian(),
            inducing_variable=Z,
        )

    def training_loop(indices: Sequence[int], num_data: int, max_iter: int) -> SVGP:
        model = get_model(num_data)
        opt = tf.optimizers.SGD(learning_rate=0.001)
        data = X[indices], Y[indices]
        for _ in range(max_iter):
            with tf.GradientTape() as tape:
                loss = model.training_loss(data)
            grads = tape.gradient(loss, model.trainable_variables)
            opt.apply_gradients(zip(grads, model.trainable_variables))
        return model

    model_1 = training_loop(indices_1, num_data=num_data1, max_iter=max_iter)
    model_2 = training_loop(indices_2, num_data=num_data2, max_iter=max_iter)
    assert _check_models_close(model_1, model_2)


# tests/gpflow/models/test_svgp_posterior.py

from itertools import product
from typing import Optional, Tuple

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import AnyNDArray
from gpflow.models.svgp import SVGP_deprecated, SVGP_with_posterior

input_dim = 7


def make_models(
    M: int = 64,
    D: int = input_dim,
    L: int = 3,
    q_diag: bool = False,
    whiten: bool = True,
    mo: Optional[Tuple[str, str]] = None,
) -> Tuple[SVGP_deprecated, SVGP_with_posterior]:
    if mo is None:
        k = gpflow.kernels.Matern52()
        Z = np.random.randn(M, D)
    else:
        kernel_type, iv_type = mo

        k_list = [gpflow.kernels.Matern52() for _ in range(L)]
        output_dim = 5
        w = tf.Variable(initial_value=np.random.rand(output_dim, L), dtype=tf.float64, name="w")
        if kernel_type == "LinearCoregionalization":
            k = gpflow.kernels.LinearCoregionalization(k_list, W=w)
        elif kernel_type == "SeparateIndependent":
            k = gpflow.kernels.SeparateIndependent(k_list)
        elif kernel_type == "SharedIndependent":
            k = gpflow.kernels.SharedIndependent(k_list[0], output_dim=L)
        else:
            raise NotImplementedError

        iv_list = [
            gpflow.inducing_variables.InducingPoints(np.random.randn(M, D)) for _ in range(L)
        ]
        if iv_type == "SeparateIndependent":
            Z = gpflow.inducing_variables.SeparateIndependentInducingVariables(iv_list)
        elif iv_type == "SharedIndependent":
            Z = gpflow.inducing_variables.SharedIndependentInducingVariables(iv_list[0])
        else:
            raise NotImplementedError

    lik = gpflow.likelihoods.Gaussian(0.1)
    q_mu = np.random.randn(M, L)
    if q_diag:
        q_sqrt: AnyNDArray = np.random.randn(M, L) ** 2
    else:
        q_sqrt = np.tril(np.random.randn(L, M, M))
    mold = SVGP_deprecated(k, lik, Z, q_diag=q_diag, q_mu=q_mu, q_sqrt=q_sqrt, whiten=whiten)
    mnew = SVGP_with_posterior(k, lik, Z, q_diag=q_diag, q_mu=q_mu, q_sqrt=q_sqrt, whiten=whiten)
    return mold, mnew


@pytest.mark.parametrize("q_diag", [True, False])
@pytest.mark.parametrize("white", [True, False])
@pytest.mark.parametrize(
    "multioutput",
    [
        None,
        *product(
            ("LinearCoregionalization", "SeparateIndependent", "SharedIndependent"),
            ("SeparateIndependent", "SharedIndependent"),
        ),
    ],
)
@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [True, False])
def test_old_vs_new_svgp(
    q_diag: bool,
    white: bool,
    multioutput: Optional[Tuple[str, str]],
    full_cov: bool,
    full_output_cov: bool,
) -> None:
    mold, mnew = make_models(q_diag=q_diag, whiten=white, mo=multioutput)

    X = np.random.randn(100, input_dim)

    mu, var = mnew.predict_f(X, full_cov=full_cov, full_output_cov=full_output_cov)
    mu2, var2 = mold.predict_f(X, full_cov=full_cov, full_output_cov=full_output_cov)
    np.testing.assert_allclose(mu, mu2)
    np.testing.assert_allclose(var, var2)


# tests/gpflow/models/test_training_mixins.py

import numpy as np
import tensorflow as tf

import gpflow
from gpflow.base import RegressionData


class DummyModel(gpflow.models.BayesianModel, gpflow.models.ExternalDataTrainingLossMixin):
    # type-ignore is because of changed method signature:
    def maximum_log_likelihood_objective(self, data: RegressionData) -> tf.Tensor:  # type: ignore[override]
        X, Y = data
        return tf.reduce_sum(X * Y)


def test_training_loss_closure_with_minibatch() -> None:
    N = 13
    B = 5
    num_batches = int(np.ceil(N / B))

    data = (np.random.randn(N, 2), np.random.randn(N, 1))
    dataset = tf.data.Dataset.from_tensor_slices(data)

    model = DummyModel()

    training_loss_full_data = model.training_loss_closure(data, compile=True)
    loss_full = training_loss_full_data()

    it = iter(dataset.batch(B))
    training_loss_minibatch = model.training_loss_closure(it, compile=True)
    batch_losses = [training_loss_minibatch() for _ in range(num_batches)]

    np.testing.assert_allclose(loss_full, np.sum(batch_losses))


# tests/gpflow/models/test_variational.py

# Copyright 2016 the GPflow authors.
#
# 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 typing import Tuple, cast

import numpy as np
import pytest
from numpy.testing import assert_allclose

import gpflow
from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import check_shapes
from gpflow.kernels import SquaredExponential
from gpflow.likelihoods import Gaussian
from tests.gpflow.kernels.reference import ref_rbf_kernel

rng = np.random.RandomState(1)

# ------------------------------------------
# Helpers
# ------------------------------------------


@check_shapes(
    "y: []",
    "K: []",
    "noise_var: []",
    "return: []",
)
def univariate_log_marginal_likelihood(
    y: AnyNDArray, K: AnyNDArray, noise_var: AnyNDArray
) -> AnyNDArray:
    return (  # type: ignore[no-any-return]
        -0.5 * y * y / (K + noise_var) - 0.5 * np.log(K + noise_var) - 0.5 * np.log(np.pi * 2.0)
    )


@check_shapes(
    "y: []",
    "K: []",
    "noise_var: []",
    "return[0]: []",
    "return[1]: []",
)
def univariate_posterior(
    y: AnyNDArray, K: AnyNDArray, noise_var: AnyNDArray
) -> Tuple[AnyNDArray, AnyNDArray]:
    mean = K * y / (K + noise_var)
    variance: AnyNDArray = K - K / (K + noise_var)
    return mean, variance


@check_shapes(
    "meanA: []",
    "meanB: []",
    "varA: []",
    "varB: []",
    "return: []",
)
def univariate_prior_KL(
    meanA: AnyNDArray, meanB: AnyNDArray, varA: AnyNDArray, varB: AnyNDArray
) -> AnyNDArray:
    # KL[ qA | qB ] = E_{qA} \log [qA / qB] where qA and qB are univariate normal distributions.
    return cast(
        AnyNDArray,
        0.5
        * (
            np.log(varB)
            - np.log(varA)
            - 1.0
            + varA / varB
            + cast(AnyNDArray, meanB - meanA) * cast(AnyNDArray, meanB - meanA) / varB
        ),
    )


@check_shapes(
    "meanA: [N, 1]",
    "covA: [N, N]",
    "meanB: [N, 1]",
    "covB: [N, N]",
    "return: [1, 1]",
)
def multivariate_prior_KL(
    meanA: AnyNDArray, covA: AnyNDArray, meanB: AnyNDArray, covB: AnyNDArray
) -> AnyNDArray:
    # KL[ qA | qB ] = E_{qA} \log [qA / qB] where qA and aB are
    # K dimensional multivariate normal distributions.
    # Analytically tractable and equal to...
    # 0.5 * (Tr(covB^{-1} covA) + (meanB - meanA)^T covB^{-1} (meanB - meanA)
    #        - K + log(det(covB)) - log (det(covA)))
    K = covA.shape[0]
    traceTerm = 0.5 * np.trace(np.linalg.solve(covB, covA))
    delta: AnyNDArray = meanB - meanA
    mahalanobisTerm = 0.5 * np.dot(delta.T, np.linalg.solve(covB, delta))
    constantTerm = -0.5 * K
    priorLogDeterminantTerm = 0.5 * np.linalg.slogdet(covB)[1]
    variationalLogDeterminantTerm = -0.5 * np.linalg.slogdet(covA)[1]
    return cast(
        AnyNDArray,
        (
            traceTerm
            + mahalanobisTerm
            + constantTerm
            + priorLogDeterminantTerm
            + variationalLogDeterminantTerm
        ),
    )


# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


class Datum:
    num_latent_gps = 1
    y_data: AnyNDArray = np.array(2.0)
    X: AnyNDArray = np.atleast_2d(np.array([0.0]))
    Y: AnyNDArray = np.atleast_2d(np.array([y_data]))
    Z = X.copy()
    zero_mean: AnyNDArray = np.array(0.0)
    K: AnyNDArray = np.array(1.0)
    noise_var: AnyNDArray = np.array(0.5)
    posterior_mean, posterior_var = univariate_posterior(y=y_data, K=K, noise_var=noise_var)
    posterior_std = np.sqrt(posterior_var)
    data = (X, Y)


class MultiDatum:
    dim = 3
    num_latent_gps = 1
    Y = rng.randn(dim, 1)
    X = rng.randn(dim, 1)
    Z = X.copy()
    noise_var = 0.5
    signal_var: AnyNDArray = np.array(1.5)
    ls: AnyNDArray = np.array(1.7)
    q_mean = rng.randn(dim, num_latent_gps)
    q_sqrt_diag = rng.rand(dim, num_latent_gps)
    q_sqrt_full: AnyNDArray = np.tril(rng.rand(dim, dim))


def test_reference_implementation_consistency() -> None:
    q_mean = rng.rand(1, 1)
    q_cov = rng.rand(1, 1)
    p_mean = rng.rand(1, 1)
    p_cov = rng.rand(1, 1)

    multivariate_KL = multivariate_prior_KL(q_mean, p_mean, q_cov, p_cov)
    univariate_KL = univariate_prior_KL(
        q_mean.squeeze(), p_mean.squeeze(), q_cov.squeeze(), p_cov.squeeze()
    )

    assert_allclose(univariate_KL, multivariate_KL.squeeze(), atol=4)


@pytest.mark.parametrize("diag", [True, False])
@pytest.mark.parametrize("whiten", [True, False])
def test_variational_univariate_prior_KL(diag: bool, whiten: bool) -> None:
    reference_kl = univariate_prior_KL(
        Datum.posterior_mean, Datum.zero_mean, Datum.posterior_var, Datum.K
    )
    q_mu: AnyNDArray = np.ones((1, Datum.num_latent_gps)) * Datum.posterior_mean
    ones = np.ones((1, Datum.num_latent_gps)) if diag else np.ones((1, 1, Datum.num_latent_gps))
    q_sqrt = ones * Datum.posterior_std
    model = gpflow.models.SVGP(
        kernel=SquaredExponential(variance=Datum.K),
        likelihood=Gaussian(),
        inducing_variable=Datum.Z,
        num_latent_gps=Datum.num_latent_gps,
        q_diag=diag,
        whiten=whiten,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )
    assert_allclose(model.prior_kl(), reference_kl, atol=4)


@pytest.mark.parametrize("diag", [True, False])
@pytest.mark.parametrize("whiten", [True, False])
def test_variational_univariate_log_likelihood(diag: bool, whiten: bool) -> None:
    # reference marginal likelihood estimate
    reference_log_marginal_likelihood = univariate_log_marginal_likelihood(
        y=Datum.y_data, K=Datum.K, noise_var=Datum.noise_var
    )
    q_mu: AnyNDArray = np.ones((1, Datum.num_latent_gps)) * Datum.posterior_mean
    ones = np.ones((1, Datum.num_latent_gps)) if diag else np.ones((1, 1, Datum.num_latent_gps))
    q_sqrt = ones * Datum.posterior_std
    model = gpflow.models.SVGP(
        kernel=SquaredExponential(variance=Datum.K),
        likelihood=Gaussian(),
        inducing_variable=Datum.Z,
        num_latent_gps=Datum.num_latent_gps,
        q_diag=diag,
        whiten=whiten,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )
    model_likelihood = model.elbo(Datum.data).numpy()
    assert_allclose(model_likelihood, reference_log_marginal_likelihood, atol=4)


@pytest.mark.parametrize("diag", [True, False])
@pytest.mark.parametrize("whiten", [True, False])
def test_variational_univariate_conditionals(diag: bool, whiten: bool) -> None:
    q_mu: AnyNDArray = np.ones((1, Datum.num_latent_gps)) * Datum.posterior_mean
    ones = np.ones((1, Datum.num_latent_gps)) if diag else np.ones((1, 1, Datum.num_latent_gps))
    q_sqrt = ones * Datum.posterior_std
    model = gpflow.models.SVGP(
        kernel=SquaredExponential(variance=Datum.K),
        likelihood=Gaussian(),
        inducing_variable=Datum.Z,
        num_latent_gps=Datum.num_latent_gps,
        q_diag=diag,
        whiten=whiten,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )

    fmean_func, fvar_func = gpflow.conditionals.conditional(
        Datum.X, Datum.Z, model.kernel, model.q_mu, q_sqrt=model.q_sqrt, white=whiten
    )
    mean_value, var_value = fmean_func[0, 0], fvar_func[0, 0]

    assert_allclose(mean_value, Datum.posterior_mean, atol=4)
    assert_allclose(var_value, Datum.posterior_var, atol=4)


@pytest.mark.parametrize("whiten", [True, False])
def test_variational_multivariate_prior_KL_full_q(whiten: bool) -> None:
    cov_q: AnyNDArray = MultiDatum.q_sqrt_full @ MultiDatum.q_sqrt_full.T
    mean_prior = np.zeros((MultiDatum.dim, 1))
    cov_prior = (
        np.eye(MultiDatum.dim)
        if whiten
        else ref_rbf_kernel(MultiDatum.X, MultiDatum.ls, MultiDatum.signal_var)
    )
    reference_kl = multivariate_prior_KL(MultiDatum.q_mean, cov_q, mean_prior, cov_prior)

    q_sqrt = MultiDatum.q_sqrt_full[None, :, :]
    model = gpflow.models.SVGP(
        kernel=SquaredExponential(variance=MultiDatum.signal_var, lengthscales=MultiDatum.ls),
        likelihood=Gaussian(MultiDatum.noise_var),
        inducing_variable=MultiDatum.Z,
        num_latent_gps=MultiDatum.num_latent_gps,
        q_diag=False,
        whiten=whiten,
        q_mu=MultiDatum.q_mean,
        q_sqrt=q_sqrt,
    )

    assert_allclose(model.prior_kl(), reference_kl, atol=4)


# tests/gpflow/models/test_vgp.py

# Copyright 2022 the GPflow authors.
#
# 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.
import numpy as np
import tensorflow as tf

import gpflow
from gpflow.config import default_float


def test_update_vgp_data() -> None:
    rng = np.random.default_rng(20220223)
    sample = lambda *shape: tf.convert_to_tensor(rng.standard_normal(shape), dtype=default_float())

    n_inputs = 2
    n_outputs = 1
    n_data_1 = 3
    n_data_2 = 2

    X_1 = tf.Variable(sample(n_data_1, n_inputs), shape=(None, n_inputs), trainable=False)
    Y_1 = tf.Variable(sample(n_data_1, n_outputs), shape=(None, n_outputs), trainable=False)

    model = gpflow.models.VGP(
        (X_1, Y_1),
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.Gaussian(),
        num_latent_gps=n_outputs,
    )
    gpflow.optimizers.Scipy().minimize(
        model.training_loss_closure(),
        variables=model.trainable_variables,
        options=dict(maxiter=25),
        compile=True,
    )

    X_test = tf.constant(tf.convert_to_tensor(X_1))
    mean_before, var_before = model.predict_f(X_test)

    X_2 = sample(n_data_2, n_inputs)
    Y_2 = sample(n_data_2, n_outputs)
    gpflow.models.vgp.update_vgp_data(
        model, (tf.concat([X_1, X_2], axis=0), tf.concat([Y_1, Y_2], axis=0))
    )

    (
        mean_after,
        var_after,
    ) = model.predict_f(X_test)

    np.testing.assert_allclose(mean_before, mean_after, atol=1e-5)
    np.testing.assert_allclose(var_before, var_after, atol=1e-6)


# tests/gpflow/models/test_vgp_posterior.py

from typing import Tuple

import numpy as np
import pytest

import gpflow
from gpflow.base import AnyNDArray, RegressionData
from gpflow.experimental.check_shapes import check_shapes
from gpflow.models.vgp import VGP_deprecated, VGP_with_posterior
from gpflow.posteriors import PrecomputeCacheType


@check_shapes(
    "regression_data[0]: [N, D]",
    "regression_data[1]: [N, P]",
)
def make_models(
    regression_data: RegressionData, likelihood: gpflow.likelihoods.Likelihood
) -> Tuple[VGP_deprecated, VGP_with_posterior]:
    """Helper function to create models"""

    k = gpflow.kernels.Matern52()
    likelihood = gpflow.likelihoods.Gaussian()

    mold = VGP_deprecated(data=regression_data, kernel=k, likelihood=likelihood)
    mnew = VGP_with_posterior(data=regression_data, kernel=k, likelihood=likelihood)
    return mold, mnew


@check_shapes(
    "return[0]: [N, D]  # X",
    "return[1]: [batch_new..., N_new, D]  # X_new",
    "return[2]: [N, P]  # Y",
)
def _get_data_for_tests() -> Tuple[AnyNDArray, AnyNDArray, AnyNDArray]:
    """Helper function to create testing data"""
    X = np.random.randn(5, 6)
    Y = np.random.randn(5, 2)
    X_new = np.random.randn(3, 10, 5, 6)
    return X, X_new, Y


@pytest.mark.parametrize(
    "likelihood", [gpflow.likelihoods.Gaussian(), gpflow.likelihoods.Exponential()]
)
@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [False])
def test_old_vs_new_gp_fused(
    likelihood: gpflow.likelihoods.Likelihood,
    full_cov: bool,
    full_output_cov: bool,
) -> None:
    X, X_new, Y = _get_data_for_tests()
    mold, mnew = make_models((X, Y), likelihood)

    mu_old, var2_old = mold.predict_f(X_new, full_cov=full_cov, full_output_cov=full_output_cov)
    mu_new_fuse, var2_new_fuse = mnew.predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )
    # check new fuse is same as old version
    np.testing.assert_allclose(mu_new_fuse, mu_old)
    np.testing.assert_allclose(var2_new_fuse, var2_old)


@pytest.mark.parametrize("cache_type", [PrecomputeCacheType.TENSOR, PrecomputeCacheType.VARIABLE])
@pytest.mark.parametrize(
    "likelihood", [gpflow.likelihoods.Gaussian(), gpflow.likelihoods.Exponential()]
)
@pytest.mark.parametrize("full_cov", [True, False])
@pytest.mark.parametrize("full_output_cov", [False])
def test_old_vs_new_with_posterior(
    cache_type: PrecomputeCacheType,
    likelihood: gpflow.likelihoods.Likelihood,
    full_cov: bool,
    full_output_cov: bool,
) -> None:
    X, X_new, Y = _get_data_for_tests()
    mold, mnew = make_models((X, Y), likelihood)

    mu_old, var2_old = mold.predict_f(X_new, full_cov=full_cov, full_output_cov=full_output_cov)
    mu_new_cache, var2_new_cache = mnew.posterior(cache_type).predict_f(
        X_new, full_cov=full_cov, full_output_cov=full_output_cov
    )

    # check new cache is same as old version
    np.testing.assert_allclose(mu_old, mu_new_cache)
    np.testing.assert_allclose(var2_old, var2_new_cache)


# tests/gpflow/optimizers/test_mcmc.py

from typing import Tuple

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.distributions import Gamma, Uniform

import gpflow
from gpflow.base import AnyNDArray, PriorOn
from gpflow.experimental.check_shapes import ShapeChecker, check_shapes
from gpflow.models import GPR
from gpflow.utilities import to_default_float

np.random.seed(1)


def build_data() -> Tuple[AnyNDArray, AnyNDArray]:
    cs = ShapeChecker().check_shape
    N = 30
    X = cs(np.random.rand(N, 1), "[N, 1]")
    Y = cs(np.sin(12 * X) + 0.66 * np.cos(25 * X) + np.random.randn(N, 1) * 0.1 + 3, "[N, 1]")
    return (X, Y)


@check_shapes(
    "data[0]: [N, 1]",
    "data[1]: [N, 1]",
)
def build_model(data: Tuple[AnyNDArray, AnyNDArray]) -> GPR:

    kernel = gpflow.kernels.Matern52(lengthscales=0.3)

    meanf = gpflow.mean_functions.Linear(1.0, 0.0)
    model = GPR(data, kernel, meanf, noise_variance=0.01)

    for p in model.parameters:
        p.prior = Gamma(to_default_float(1.0), to_default_float(1.0))

    return model


@check_shapes(
    "data[0]: [N, 1]",
    "data[1]: [N, 1]",
)
def build_model_with_uniform_prior_no_transforms(
    data: Tuple[AnyNDArray, AnyNDArray], prior_on: PriorOn, prior_width: float
) -> GPR:
    def parameter(value: tf.Tensor) -> gpflow.Parameter:
        low_value = -100
        high_value = low_value + prior_width
        prior = Uniform(low=np.float64(low_value), high=np.float64(high_value))
        return gpflow.Parameter(
            value, transform=tfp.bijectors.Identity(), prior=prior, prior_on=prior_on
        )

    k = gpflow.kernels.Matern52(lengthscales=0.3)
    k.variance = parameter(k.variance)
    k.lengthscales = parameter(k.lengthscales)

    mf = gpflow.mean_functions.Linear(1.0, 0.0)
    mf.A = parameter(mf.A)
    mf.b = parameter(mf.b)
    m = GPR(data, k, mf, noise_variance=0.01)
    m.likelihood.variance = parameter(m.likelihood.variance)
    return m


def test_mcmc_helper_parameters() -> None:
    data = build_data()
    model = build_model(data)

    hmc_helper = gpflow.optimizers.SamplingHelper(
        model.log_posterior_density, model.trainable_parameters
    )

    for i in range(len(model.trainable_parameters)):
        assert model.trainable_parameters[i].shape == hmc_helper.current_state[i].shape
        assert model.trainable_parameters[i] == hmc_helper._parameters[i]
        assert model.trainable_parameters[i].unconstrained_variable == hmc_helper.current_state[i]


def test_mcmc_helper_target_function_constrained() -> None:
    """Set up priors on the model parameters such that we can
    readily compute their expected values."""
    config = gpflow.config.Config(positive_bijector="exp")
    with gpflow.config.as_context(config):
        data = build_data()
        model = build_model(data)

    prior_width = 200.0

    hmc_helper = gpflow.optimizers.SamplingHelper(
        model.log_posterior_density, model.trainable_parameters
    )
    target_log_prob_fn = hmc_helper.target_log_prob_fn

    # Priors which are set on the constrained space
    expected_log_prior = 0.0
    for param in model.trainable_parameters:
        if param.numpy() < 1e-3:
            # Avoid values which would be pathological for the Exp transform
            param.assign(1.0)

        low_value = -100
        high_value = low_value + prior_width

        param.prior = Uniform(low=np.float64(low_value), high=np.float64(high_value))
        param.prior_on = PriorOn.CONSTRAINED

        prior_density_on_constrained = 1 / prior_width
        prior_density_on_unconstrained = prior_density_on_constrained * param.numpy()

        expected_log_prior += np.log(prior_density_on_unconstrained)

    log_marginal_likelihood = model.log_marginal_likelihood().numpy()
    expected_log_prob = log_marginal_likelihood + expected_log_prior

    np.testing.assert_allclose(target_log_prob_fn(), expected_log_prob, rtol=1e-6)


def test_mcmc_helper_target_function_unconstrained() -> None:
    """
    Verifies the objective for a set of priors which are defined on the unconstrained space.
    """
    # Set up priors on the model parameters such that we can readily compute their expected values.
    expected_log_prior = 0.0
    prior_width = 200.0

    data = build_data()
    model = build_model_with_uniform_prior_no_transforms(data, PriorOn.UNCONSTRAINED, prior_width)

    hmc_helper = gpflow.optimizers.SamplingHelper(
        model.log_posterior_density, model.trainable_parameters
    )

    for _ in model.trainable_parameters:
        prior_density = 1 / prior_width
        expected_log_prior += np.log(prior_density)

    target_log_prob_fn = hmc_helper.target_log_prob_fn
    expected_log_prob = model.log_marginal_likelihood().numpy() + expected_log_prior

    np.testing.assert_allclose(target_log_prob_fn(), expected_log_prob)


@pytest.mark.parametrize("prior_on", [PriorOn.CONSTRAINED, PriorOn.UNCONSTRAINED])
def test_mcmc_helper_target_function_no_transforms(prior_on: PriorOn) -> None:
    """Verifies the objective for a set of priors where no transforms are set."""
    expected_log_prior = 0.0
    prior_width = 200.0

    data = build_data()
    model = build_model_with_uniform_prior_no_transforms(data, prior_on, prior_width)

    hmc_helper = gpflow.optimizers.SamplingHelper(
        model.log_posterior_density, model.trainable_parameters
    )

    for _ in model.trainable_parameters:
        prior_density = 1 / prior_width
        expected_log_prior += np.log(prior_density)

    log_marginal_likelihood = model.log_marginal_likelihood().numpy()
    expected_log_prob = log_marginal_likelihood + expected_log_prior
    target_log_prob_fn = hmc_helper.target_log_prob_fn

    np.testing.assert_allclose(target_log_prob_fn(), expected_log_prob)

    # Test the wrapped closure
    log_prob, grad_fn = target_log_prob_fn.__original_wrapped__()  # type: ignore[attr-defined]
    grad, nones = grad_fn(1, [None] * len(model.trainable_parameters))
    assert len(grad) == len(model.trainable_parameters)
    assert nones == [None] * len(model.trainable_parameters)


def test_mcmc_sampler_integration() -> None:
    data = build_data()
    model = build_model(data)

    hmc_helper = gpflow.optimizers.SamplingHelper(
        model.log_posterior_density, model.trainable_parameters
    )

    hmc = tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=hmc_helper.target_log_prob_fn,
        num_leapfrog_steps=2,
        step_size=0.01,
    )

    adaptive_hmc = tfp.mcmc.SimpleStepSizeAdaptation(
        hmc,
        num_adaptation_steps=2,
        target_accept_prob=gpflow.utilities.to_default_float(0.75),
        adaptation_rate=0.1,
    )

    num_samples = 5

    @tf.function
    def run_chain_fn() -> tfp.mcmc.StatesAndTrace:
        return tfp.mcmc.sample_chain(
            num_results=num_samples,
            num_burnin_steps=2,
            current_state=hmc_helper.current_state,
            kernel=adaptive_hmc,
            trace_fn=lambda _, pkr: pkr.inner_results.is_accepted,
        )

    samples, _ = run_chain_fn()

    assert len(samples) == len(model.trainable_parameters)
    parameter_samples = hmc_helper.convert_to_constrained_values(samples)
    assert len(parameter_samples) == len(samples)

    for i in range(len(model.trainable_parameters)):
        assert len(samples[i]) == num_samples
        assert hmc_helper.current_state[i].numpy() == samples[i][-1]
        assert hmc_helper._parameters[i].numpy() == parameter_samples[i][-1]


def test_helper_with_variables_fails() -> None:
    variable = tf.Variable(0.1)
    with pytest.raises(
        ValueError, match=r"`parameters` should only contain gpflow.Parameter objects with priors"
    ):
        gpflow.optimizers.SamplingHelper(lambda: variable ** 2, (variable,))


# tests/gpflow/optimizers/test_natural_gradient.py

from typing import Any, Optional, Tuple

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow import set_trainable
from gpflow.config import default_float
from gpflow.experimental.check_shapes import check_shapes
from gpflow.kernels import Kernel, SquaredExponential
from gpflow.likelihoods import Gaussian, Likelihood
from gpflow.models import GPR, SGPR, SVGP, VGP, BayesianModel
from gpflow.optimizers import NaturalGradient


class Setup:
    N, M, D = 4, 3, 2
    likelihood_variance = 0.1
    rng = np.random.RandomState(42)
    X = rng.randn(N, D)
    Y = rng.randn(N, 1)
    Z = rng.randn(M, D)


@pytest.fixture
@check_shapes(
    "return[0]: [N, D]",
    "return[1]: [N, 1]",
)
def data() -> Tuple[tf.Tensor, tf.Tensor]:
    X = tf.convert_to_tensor(Setup.X, dtype=default_float())
    Y = tf.convert_to_tensor(Setup.Y, dtype=default_float())
    return (X, Y)


@pytest.fixture
@check_shapes(
    "return: [N, D]",
)
def inducing_variable() -> tf.Tensor:
    Z = tf.convert_to_tensor(Setup.Z, dtype=default_float())
    return Z


@pytest.fixture
def kernel() -> Kernel:
    return SquaredExponential()


@pytest.fixture
def likelihood() -> Likelihood:
    return Gaussian(variance=Setup.likelihood_variance)


@pytest.fixture
def gpr_and_vgp(
    data: Tuple[tf.Tensor, tf.Tensor], kernel: Kernel, likelihood: Likelihood
) -> Tuple[GPR, VGP]:
    vgp = VGP(data, kernel, likelihood)
    gpr = GPR(data, kernel)
    gpr.likelihood.variance.assign(likelihood.variance)
    set_trainable(vgp, False)
    set_trainable(vgp.q_mu, True)
    set_trainable(vgp.q_sqrt, True)
    return gpr, vgp


@pytest.fixture
def sgpr_and_svgp(
    data: Tuple[tf.Tensor, tf.Tensor],
    inducing_variable: tf.Tensor,
    kernel: Kernel,
    likelihood: Likelihood,
) -> Tuple[SGPR, SVGP]:
    svgp = SVGP(kernel, likelihood, inducing_variable)
    sgpr = SGPR(data, kernel, inducing_variable=inducing_variable)
    sgpr.likelihood.variance.assign(Setup.likelihood_variance)
    set_trainable(svgp, False)
    set_trainable(svgp.q_mu, True)
    set_trainable(svgp.q_sqrt, True)
    return sgpr, svgp


@check_shapes(
    "value1: [broadcast batch...]",
    "value2: [broadcast batch...]",
)
def assert_different(value1: tf.Tensor, value2: tf.Tensor, rtol: float = 0.07) -> None:
    """ assert relative difference > rtol """
    relative_difference = (value1 - value2) / (value1 + value2)
    assert np.abs(relative_difference) > rtol


def assert_gpr_vs_vgp(
    m1: BayesianModel,
    m2: BayesianModel,
    gamma: float = 1.0,
    maxiter: int = 1,
    xi_transform: Optional[gpflow.optimizers.natgrad.XiTransform] = None,
) -> None:
    assert maxiter >= 1

    m1_ll_before = m1.training_loss()
    m2_ll_before = m2.training_loss()

    assert_different(m2_ll_before, m1_ll_before)

    params: Tuple[Any, ...] = (m2.q_mu, m2.q_sqrt)
    if xi_transform is not None:
        params += (xi_transform,)

    opt = NaturalGradient(gamma)

    @tf.function
    def minimize_step() -> None:
        opt.minimize(m2.training_loss, var_list=[params])  # type: ignore[list-item]

    for _ in range(maxiter):
        minimize_step()

    m1_ll_after = m1.training_loss()
    m2_ll_after = m2.training_loss()

    np.testing.assert_allclose(m1_ll_after, m2_ll_after, atol=1e-4)


def assert_sgpr_vs_svgp(
    m1: BayesianModel,
    m2: BayesianModel,
) -> None:
    data = m1.data

    m1_ll_before = m1.training_loss()
    m2_ll_before = m2.training_loss(data)

    assert_different(m2_ll_before, m1_ll_before)

    params = [(m2.q_mu, m2.q_sqrt)]
    opt = NaturalGradient(1.0)
    opt.minimize(m2.training_loss_closure(data), var_list=params)

    m1_ll_after = m1.training_loss()
    m2_ll_after = m2.training_loss(data)

    np.testing.assert_allclose(m1_ll_after, m2_ll_after, atol=1e-4)


def test_vgp_vs_gpr(gpr_and_vgp: Tuple[GPR, VGP]) -> None:
    """
    With a Gaussian likelihood the Gaussian variational (VGP) model should be equivalent to the
    exact regression model (GPR) after a single nat grad step of size 1
    """
    gpr, vgp = gpr_and_vgp
    assert_gpr_vs_vgp(gpr, vgp)


def test_small_q_sqrt_handeled_correctly(
    gpr_and_vgp: Tuple[GPR, VGP], data: Tuple[tf.Tensor, tf.Tensor]
) -> None:
    """
    This is an extra test to make sure things still work when q_sqrt is small.
    This was breaking (#767)
    """
    gpr, vgp = gpr_and_vgp
    vgp.q_mu.assign(np.random.randn(data[0].shape[0], 1))
    vgp.q_sqrt.assign(np.eye(data[0].shape[0])[None, :, :] * 1e-3)
    assert_gpr_vs_vgp(gpr, vgp)


def test_svgp_vs_sgpr(sgpr_and_svgp: Tuple[SGPR, SVGP]) -> None:
    """
    With a Gaussian likelihood the sparse Gaussian variational (SVGP) model
    should be equivalent to the analytically optimial sparse regression model (SGPR)
    after a single nat grad step of size 1.0
    """
    sgpr, svgp = sgpr_and_svgp
    assert_sgpr_vs_svgp(sgpr, svgp)


class XiEta(gpflow.optimizers.XiTransform):
    @staticmethod
    @check_shapes(
        "mean: [N, D]",
        "varsqrt: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def meanvarsqrt_to_xi(mean: tf.Tensor, varsqrt: tf.Tensor) -> tf.Tensor:
        return gpflow.optimizers.natgrad.meanvarsqrt_to_expectation(mean, varsqrt)

    @staticmethod
    @check_shapes(
        "xi1: [N, D]",
        "xi2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def xi_to_meanvarsqrt(xi1: tf.Tensor, xi2: tf.Tensor) -> tf.Tensor:
        return gpflow.optimizers.natgrad.expectation_to_meanvarsqrt(xi1, xi2)

    @staticmethod
    @check_shapes(
        "nat1: [N, D]",
        "nat2: [D, N, N]",
        "return[0]: [N, D]",
        "return[1]: [D, N, N]",
    )
    def naturals_to_xi(nat1: tf.Tensor, nat2: tf.Tensor) -> tf.Tensor:
        return gpflow.optimizers.natgrad.natural_to_expectation(nat1, nat2)


@pytest.mark.parametrize("xi_transform", [gpflow.optimizers.XiSqrtMeanVar(), XiEta()])
def test_xi_transform_vgp_vs_gpr(
    gpr_and_vgp: Tuple[GPR, VGP], xi_transform: gpflow.optimizers.XiTransform
) -> None:
    """
    With other transforms the solution is not given in a single step, but it should still give the same answer
    after a number of smaller steps.
    """
    gpr, vgp = gpr_and_vgp
    assert_gpr_vs_vgp(gpr, vgp, gamma=0.01, xi_transform=xi_transform, maxiter=500)


# tests/gpflow/optimizers/test_scipy.py

# Copyright 2019 the GPflow authors.
#
# 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.

import warnings

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import AnyNDArray
from gpflow.config import default_float
from gpflow.models import GPR, GPModel

rng = np.random.RandomState(0)


class Datum:
    X: AnyNDArray = rng.rand(20, 1) * 10
    Y = np.sin(X) + 0.9 * np.cos(X * 1.6) + rng.randn(*X.shape) * 0.8
    Y = np.tile(Y, 2)  # two identical columns
    Xtest: AnyNDArray = rng.rand(10, 1) * 10
    data = (X, Y)


def _create_full_gp_model() -> GPModel:
    """
    GP Regression
    """
    return GPR(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        mean_function=gpflow.mean_functions.Constant(),
    )


def test_scipy_jit() -> None:
    m1 = _create_full_gp_model()
    m2 = _create_full_gp_model()

    opt1 = gpflow.optimizers.Scipy()
    opt2 = gpflow.optimizers.Scipy()

    opt1.minimize(
        m1.training_loss,
        variables=m1.trainable_variables,
        options=dict(maxiter=50),
        compile=False,
    )
    opt2.minimize(
        m2.training_loss,
        variables=m2.trainable_variables,
        options=dict(maxiter=50),
        compile=True,
    )

    def get_values(model: GPModel) -> AnyNDArray:
        return np.array([var.numpy().squeeze() for var in model.trainable_variables])

    np.testing.assert_allclose(get_values(m1), get_values(m2), rtol=1e-13, atol=1e-13)


@pytest.mark.parametrize("compile", [True, False])
def test_scipy__optimal(compile: bool) -> None:
    target1 = [0.2, 0.8]
    target2 = [0.6]
    v1 = tf.Variable([0.5, 0.5], dtype=default_float())
    v2 = tf.Variable([0.5], dtype=default_float())
    compilation_count = 0

    def f() -> tf.Tensor:
        nonlocal compilation_count
        compilation_count += 1
        return tf.reduce_sum((target1 - v1) ** 2) + tf.reduce_sum((target2 - v2) ** 2)

    opt = gpflow.optimizers.Scipy()
    result = opt.minimize(f, [v1, v2], compile=compile)

    if compile:
        assert 1 == compilation_count
    else:
        assert 1 < compilation_count
    assert result.success
    np.testing.assert_allclose(target1 + target2, result.x)
    np.testing.assert_allclose(target1, v1)
    np.testing.assert_allclose(target2, v2)


@pytest.mark.parametrize("compile", [True, False])
def test_scipy__partially_disconnected_variable(compile: bool) -> None:
    target1 = 0.2
    target2 = 0.6
    v1 = tf.Variable([0.5, 0.5], dtype=default_float())
    v2 = tf.Variable(0.5, dtype=default_float())

    def f() -> tf.Tensor:
        # v1[1] not used.
        v10 = v1[0]
        return (target1 - v10) ** 2 + (target2 - v2) ** 2

    opt = gpflow.optimizers.Scipy()
    result = opt.minimize(f, [v1, v2], compile=compile)

    assert result.success
    np.testing.assert_allclose([target1, 0.5, target2], result.x)
    np.testing.assert_allclose([target1, 0.5], v1)
    np.testing.assert_allclose(target2, v2)


@pytest.mark.parametrize("compile", [True, False])
@pytest.mark.parametrize("allow_unused_variables", [True, False])
def test_scipy__disconnected_variable(compile: bool, allow_unused_variables: bool) -> None:
    target1 = [0.2, 0.8]
    v1 = tf.Variable([0.5, 0.5], dtype=default_float(), name="v1")
    v2 = tf.Variable([0.5], dtype=default_float(), name="v2")

    def f() -> tf.Tensor:
        # v2 not used.
        return tf.reduce_sum((target1 - v1) ** 2)

    opt = gpflow.optimizers.Scipy()

    if allow_unused_variables:
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            result = opt.minimize(
                f, [v1, v2], compile=compile, allow_unused_variables=allow_unused_variables
            )

        (warning,) = w
        message = warning.message
        assert isinstance(message, Warning)
        msg = message.args[0]
        assert v2.name in msg

        assert result.success
        np.testing.assert_allclose(target1 + [0.5], result.x)
        np.testing.assert_allclose(target1, v1)
        np.testing.assert_allclose([0.5], v2)
    else:
        with pytest.raises(ValueError, match=v2.name):
            opt.minimize(f, [v1, v2], allow_unused_variables=allow_unused_variables)


# tests/gpflow/posteriors/conftest.py

#  Copyright 2022 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 inspect import isabstract
from typing import DefaultDict, Iterable, Set, Type

import pytest

import gpflow.ci_utils
from gpflow.posteriors import AbstractPosterior


@pytest.fixture(name="tested_posteriors", scope="package")
def _tested_posteriors() -> DefaultDict[str, Set[Type[AbstractPosterior]]]:
    return DefaultDict(set)


@pytest.fixture(scope="package", autouse=True)
def _ensure_all_posteriors_are_tested_fixture(
    tested_posteriors: DefaultDict[str, Set[Type[AbstractPosterior]]]
) -> Iterable[None]:
    """
    This fixture ensures that all concrete posteriors have unit tests which compare the predictions
    from the fused and precomputed code paths. When adding a new concrete posterior class to
    GPFlow, ensure that it is also tested in this manner.

    This autouse, package scoped fixture will always be executed when tests in this package are run.
    """
    # Code here will be executed before any of the tests in this package.

    yield  # Run tests in this package.

    # Code here will be executed after all of the tests in this package.

    available_posteriors = list(gpflow.ci_utils.subclasses(AbstractPosterior))
    concrete_posteriors = set([k for k in available_posteriors if not isabstract(k)])

    messages = []
    for key, key_tested_posteriors in tested_posteriors.items():
        untested_posteriors = concrete_posteriors - key_tested_posteriors
        if untested_posteriors:
            messages.append(
                f"For key '{key}' no tests have been registered for the following posteriors: {untested_posteriors}."
            )

    if messages:
        raise AssertionError("\n".join(messages))


# tests/gpflow/posteriors/test_bo_integration.py

#  Copyright 2022 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 typing import (
    Any,
    Callable,
    DefaultDict,
    Generic,
    Iterator,
    List,
    Set,
    Tuple,
    Type,
    TypeVar,
    cast,
)

import numpy as np
import pytest
import tensorflow as tf
from _pytest.fixtures import SubRequest

import gpflow
from gpflow.base import RegressionData
from gpflow.config import default_float
from gpflow.experimental.check_shapes import check_shapes
from gpflow.inducing_variables import InducingPoints, InducingVariables
from gpflow.kernels import Kernel, Matern52
from gpflow.likelihoods import Exponential, Likelihood
from gpflow.models import GPR, SGPR, SVGP, VGP, GPModel, training_loss_closure
from gpflow.models.vgp import update_vgp_data
from gpflow.posteriors import AbstractPosterior, PrecomputeCacheType

_M = TypeVar("_M", bound=GPModel, covariant=True)
_CreateModel = Callable[[RegressionData], _M]


# I'd like to make this a `dataclass`, but mypy get confused about `create_model` being a function
# member, but that doesn't take `self`.
class _ModelFactory(Generic[_M]):
    def __init__(
        self,
        create_model: _CreateModel[_M],
        multi_output: bool,
        atol: float,
        rtol: float,
    ) -> None:
        self.create_model = create_model
        self.multi_output = multi_output
        self.atol = atol
        self.rtol = rtol


_MODEL_FACTORIES: List[_ModelFactory[Any]] = []

# This exists to make it easy to disable tf.function, for debugging.
_COMPILE = True
_MAXITER = 10
_DEFAULT_ATOL = 1e-10
_DEFAULT_RTOL = 1e-7


@pytest.fixture(name="register_posterior_bo_integration_test")
def _register_posterior_bo_integration_test(
    request: SubRequest,
    tested_posteriors: DefaultDict[str, Set[Type[AbstractPosterior]]],
) -> Callable[[AbstractPosterior], None]:
    def _register_posterior(posterior: AbstractPosterior) -> None:
        tested_posteriors[request.function.__name__].add(posterior.__class__)

    return _register_posterior


def model_factory(
    multi_output: bool = False, atol: float = _DEFAULT_ATOL, rtol: float = _DEFAULT_RTOL
) -> Callable[[_CreateModel[_M]], _ModelFactory[_M]]:
    """ Decorator for adding a function to the `_MODEL_FACTORIES` list. """

    def register(create_model: _CreateModel[_M]) -> _ModelFactory[_M]:
        model_factory = _ModelFactory(
            create_model,
            multi_output,
            atol,
            rtol,
        )
        _MODEL_FACTORIES.append(model_factory)
        return model_factory

    return register


def create_kernel() -> Kernel:
    return Matern52()


def create_likelihood() -> Likelihood:
    return Exponential()


def create_inducing_points(data: RegressionData) -> InducingPoints:
    n_features = data[0].shape[1]
    n_inducing_points = 5
    rng = np.random.default_rng(20220208)
    Z = tf.constant(rng.random((n_inducing_points, n_features)))
    return InducingPoints(Z)


@check_shapes(
    "inducing_variable: [M, D, L]",
    "return[1]: [M_x_L, P]",
    "return[2]: [M_x_L, P]",
)
def create_q(
    inducing_variable: InducingVariables, *, row_scale: int = 1, column_scale: int = 1
) -> Tuple[bool, tf.Tensor, tf.Tensor]:
    n_inducing_points = inducing_variable.num_inducing
    rng = np.random.default_rng(20220133)
    q_diag = True
    q_mu = tf.constant(rng.random((row_scale * n_inducing_points, column_scale)))
    q_sqrt = tf.constant(rng.random((row_scale * n_inducing_points, column_scale))) ** 2
    return q_diag, q_mu, q_sqrt


@model_factory(rtol=1e-3)
def create_gpr(data: RegressionData) -> GPR:
    return GPR(data=data, kernel=create_kernel())


@model_factory(rtol=1e-4)
def create_sgpr(data: RegressionData) -> SGPR:
    return SGPR(data=data, kernel=create_kernel(), inducing_variable=create_inducing_points(data))


@model_factory(rtol=5e-3)
def create_vgp(data: RegressionData) -> VGP:
    return VGP(data=data, kernel=create_kernel(), likelihood=create_likelihood())


@model_factory()
def create_svgp__independent_single_output(data: RegressionData) -> SVGP:
    inducing_variable = create_inducing_points(data)
    q_diag, q_mu, q_sqrt = create_q(inducing_variable)
    return SVGP(
        kernel=create_kernel(),
        likelihood=create_likelihood(),
        inducing_variable=inducing_variable,
        q_diag=q_diag,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )


@model_factory(multi_output=True)
def create_svgp__fully_correlated_multi_output(data: RegressionData) -> SVGP:
    n_outputs = data[1].shape[1]
    kernel = gpflow.kernels.SharedIndependent(create_kernel(), output_dim=n_outputs)
    inducing_variable = create_inducing_points(data)
    q_diag, q_mu, q_sqrt = create_q(inducing_variable, row_scale=n_outputs)
    return SVGP(
        kernel=kernel,
        likelihood=create_likelihood(),
        inducing_variable=inducing_variable,
        q_diag=q_diag,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )


@model_factory(multi_output=True)
def create_svgp__independent_multi_output(data: RegressionData) -> SVGP:
    n_outputs = data[1].shape[1]
    kernel = gpflow.kernels.SharedIndependent(create_kernel(), output_dim=n_outputs)
    inducing_variable = gpflow.inducing_variables.SharedIndependentInducingVariables(
        create_inducing_points(data)
    )
    q_diag, q_mu, q_sqrt = create_q(inducing_variable, column_scale=n_outputs)
    return SVGP(
        kernel=kernel,
        likelihood=create_likelihood(),
        inducing_variable=inducing_variable,
        q_diag=q_diag,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )


@model_factory(multi_output=True)
def create_svgp__fallback_independent_latent_posterior(data: RegressionData) -> SVGP:
    n_outputs = data[1].shape[1]
    rng = np.random.default_rng(20220131)
    kernel = gpflow.kernels.LinearCoregionalization(
        [create_kernel()],
        W=tf.constant(rng.standard_normal((n_outputs, 1))),
    )
    inducing_variable = gpflow.inducing_variables.FallbackSeparateIndependentInducingVariables(
        [create_inducing_points(data)]
    )
    q_diag, q_mu, q_sqrt = create_q(inducing_variable)
    return SVGP(
        kernel=kernel,
        likelihood=create_likelihood(),
        inducing_variable=inducing_variable,
        q_diag=q_diag,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )


@model_factory(multi_output=True)
def create_svgp__linear_coregionalization(data: RegressionData) -> SVGP:
    n_outputs = data[1].shape[1]
    rng = np.random.default_rng(20220131)
    kernel = gpflow.kernels.LinearCoregionalization(
        [create_kernel()], W=tf.constant(rng.standard_normal((n_outputs, 1)))
    )
    inducing_variable = gpflow.inducing_variables.SharedIndependentInducingVariables(
        create_inducing_points(data)
    )
    q_diag, q_mu, q_sqrt = create_q(inducing_variable)
    return SVGP(
        kernel=kernel,
        likelihood=create_likelihood(),
        inducing_variable=inducing_variable,
        q_diag=q_diag,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
    )


@pytest.fixture(params=_MODEL_FACTORIES)
def _model_factory(request: SubRequest) -> _ModelFactory[Any]:
    return cast(_ModelFactory[Any], request.param)


@pytest.fixture
@check_shapes(
    "return: [Q, D]",
)
def _f_minimum(_model_factory: _ModelFactory[Any]) -> tf.Tensor:
    return (
        tf.constant(
            [
                [0.2, 0.4],
                [0.4, 0.6],
                [0.6, 0.8],
            ],
            dtype=default_float(),
        )
        if _model_factory.multi_output
        else tf.constant([[0.3, 0.5]], dtype=default_float())
    )


@pytest.fixture
def _f(_f_minimum: tf.Tensor) -> Callable[[tf.Tensor], tf.Tensor]:
    @check_shapes(
        "X: [N, D]",
        "return: [N, Q]",
    )
    def f(X: tf.Tensor) -> tf.Tensor:
        err = X[:, None, :] - _f_minimum[None, :, :]
        err_sq = err ** 2
        return tf.reduce_sum(err_sq, axis=-1)

    return f


@pytest.fixture
@check_shapes(
    "_f_minimum: [Q, D]",
    "return[0]: [N, D]",
    "return[1]: [N, Q]",
)
def _data(
    _f: Callable[[tf.Tensor], tf.Tensor], _f_minimum: tf.Tensor
) -> Tuple[tf.Variable, tf.Variable]:
    n_initial_data = 3
    n_outputs, n_features = _f_minimum.shape

    rng = np.random.default_rng(20220126)
    X = tf.Variable(
        rng.random((n_initial_data, n_features)),
        shape=[None, n_features],
        dtype=default_float(),
        trainable=False,
    )
    Y = tf.Variable(
        _f(X),
        shape=[None, n_outputs],
        dtype=default_float(),
        trainable=False,
    )

    return X, Y


@pytest.fixture
def _extend_data(
    _data: Tuple[tf.Variable, tf.Variable], _f: Callable[[tf.Tensor], tf.Tensor]
) -> Callable[[GPModel], Iterator[int]]:
    n_iterations = 3
    rng = np.random.default_rng(20220127)
    X, Y = _data
    n_features = X.shape[1]

    def iterate(model: GPModel) -> Iterator[int]:
        for i in range(n_iterations):
            X_new = tf.constant(rng.random((1, n_features)))
            Y_new = _f(X_new)
            X_i = tf.concat([X, X_new], axis=0)
            Y_i = tf.concat([Y, Y_new], axis=0)

            if isinstance(model, VGP):
                update_vgp_data(model, (X_i, Y_i))
            else:
                X.assign(X_i)
                Y.assign(Y_i)
            yield i

    return iterate


@pytest.fixture
@check_shapes(
    "_data[0]: [N, D]",
    "_data[1]: [N, Q]",
    "return: [M, D]",
)
def _X_new(_data: Tuple[tf.Variable, tf.Variable]) -> tf.Tensor:
    rng = np.random.default_rng(20220128)
    X, _Y = _data
    n_features = X.shape[1]
    return tf.constant(rng.random((3, n_features)))


@pytest.fixture
def _optimize(_data: Tuple[tf.Variable, tf.Variable]) -> Callable[[GPModel], None]:
    def optimize(model: GPModel) -> None:
        gpflow.optimizers.Scipy().minimize(
            training_loss_closure(model, _data, compile=_COMPILE),
            variables=model.trainable_variables,
            options=dict(maxiter=_MAXITER),
            method="BFGS",
            compile=_COMPILE,
        )

    return optimize


def test_posterior_bo_integration__predict_f(
    register_posterior_bo_integration_test: Callable[[AbstractPosterior], None],
    _model_factory: _ModelFactory[Any],
    _data: Tuple[tf.Variable, tf.Variable],
    _extend_data: Callable[[GPModel], Iterator[int]],
    _X_new: tf.Tensor,
) -> None:
    """
    Check that data added incrementally is correctly reflected in `predict_f`.
    """
    _X, Y = _data
    n_rows_new = _X_new.shape[0]
    n_outputs = Y.shape[1]

    model = _model_factory.create_model(_data)
    posterior = model.posterior(PrecomputeCacheType.VARIABLE)
    register_posterior_bo_integration_test(posterior)
    if _COMPILE:
        predict_f = tf.function(lambda Xnew: posterior.predict_f(Xnew))
    else:
        predict_f = posterior.predict_f

    for _ in _extend_data(model):
        posterior.update_cache()
        compiled_mean, compiled_var = predict_f(_X_new)

        np.testing.assert_equal((n_rows_new, n_outputs), compiled_mean.shape)
        np.testing.assert_equal((n_rows_new, n_outputs), compiled_var.shape)

        eager_model = _model_factory.create_model(_data)
        eager_mean, eager_var = eager_model.predict_f(_X_new)

        np.testing.assert_allclose(
            eager_mean, compiled_mean, rtol=_model_factory.rtol, atol=_model_factory.atol
        )
        np.testing.assert_allclose(
            eager_var, compiled_var, rtol=_model_factory.rtol, atol=_model_factory.atol
        )


def test_posterior_bo_integration__optimization(
    register_posterior_bo_integration_test: Callable[[AbstractPosterior], None],
    _model_factory: _ModelFactory[Any],
    _data: Tuple[tf.Variable, tf.Variable],
    _extend_data: Callable[[GPModel], Iterator[int]],
    _X_new: tf.Tensor,
    _optimize: Callable[[GPModel], None],
) -> None:
    """
    Check that data added incrementally is considered when optimizing a model.
    """
    _X, Y = _data
    n_rows_new = _X_new.shape[0]
    n_outputs = Y.shape[1]

    model = _model_factory.create_model(_data)
    posterior = model.posterior(PrecomputeCacheType.VARIABLE)
    register_posterior_bo_integration_test(posterior)
    if _COMPILE:
        predict_f = tf.function(lambda Xnew: posterior.predict_f(Xnew))
    else:
        predict_f = posterior.predict_f

    # Add all the data first, and then `optimize`, so that both models are optimized the same number
    # of times and with the same data, so they converge to the same result.

    for _ in _extend_data(model):
        pass

    _optimize(model)
    posterior.update_cache()
    compiled_mean, compiled_var = predict_f(_X_new)

    np.testing.assert_equal((n_rows_new, n_outputs), compiled_mean.shape)
    np.testing.assert_equal((n_rows_new, n_outputs), compiled_var.shape)

    eager_model = _model_factory.create_model(_data)
    _optimize(eager_model)
    eager_mean, eager_var = eager_model.predict_f(_X_new)

    np.testing.assert_allclose(
        eager_mean, compiled_mean, rtol=_model_factory.rtol, atol=_model_factory.atol
    )
    np.testing.assert_allclose(
        eager_var, compiled_var, rtol=_model_factory.rtol, atol=_model_factory.atol
    )


# tests/gpflow/posteriors/test_posteriors.py

#  Copyright 2021 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 typing import Any, Callable, DefaultDict, Optional, Set, Type, cast

import numpy as np
import pytest
import tensorflow as tf
from _pytest.fixtures import SubRequest

import gpflow
import gpflow.ci_utils
from gpflow.base import TensorType
from gpflow.conditionals import conditional
from gpflow.experimental.check_shapes import check_shapes
from gpflow.inducing_variables import InducingPoints, InducingVariables
from gpflow.kernels import Kernel
from gpflow.likelihoods import Gaussian
from gpflow.mean_functions import Zero
from gpflow.models.util import inducingpoint_wrapper
from gpflow.posteriors import (
    AbstractPosterior,
    FallbackIndependentLatentPosterior,
    FullyCorrelatedPosterior,
    GPRPosterior,
    IndependentPosteriorMultiOutput,
    IndependentPosteriorSingleOutput,
    LinearCoregionalizationPosterior,
    PrecomputeCacheType,
    SGPRPosterior,
    VGPPosterior,
    create_posterior,
)

INPUT_DIMS = 2
NUM_INDUCING_POINTS = 3


# `PosteriorType` really should be something like `Type[AbstractPosterior]`, except mypy doesn't
# allow passing abstract classes to functions. See: https://github.com/python/mypy/issues/4717
PosteriorType = Type[Any]
RegisterPosterior = Callable[[AbstractPosterior, PosteriorType], None]


@pytest.fixture(name="register_posterior_test")
def _register_posterior_test_fixture(
    tested_posteriors: DefaultDict[str, Set[Type[AbstractPosterior]]]
) -> RegisterPosterior:
    def _verify_and_register_posterior_test(
        posterior: AbstractPosterior, expected_posterior_class: Type[AbstractPosterior]
    ) -> None:
        assert isinstance(posterior, expected_posterior_class)
        tested_posteriors["test_posteriors.py"].add(expected_posterior_class)

    return _verify_and_register_posterior_test


QSqrtFactory = Callable[[int, int], Optional[TensorType]]


@pytest.fixture(name="q_sqrt_factory", params=[0, 1, 2])
def _q_sqrt_factory_fixture(request: SubRequest) -> QSqrtFactory:
    """
    When upgrading to Python 3.10, this can be replaced with a match-case statement.
    """
    if request.param == 0:

        def fn_0(n_inducing_points: int, num_latent_gps: int) -> Optional[TensorType]:
            return None

        return fn_0
    elif request.param == 1:

        def fn_1(n_inducing_points: int, num_latent_gps: int) -> Optional[TensorType]:
            # qsqrt: [M, L]
            return tf.random.normal((n_inducing_points, num_latent_gps), dtype=tf.float64) ** 2

        return fn_1
    elif request.param == 2:

        def fn_2(n_inducing_points: int, num_latent_gps: int) -> Optional[TensorType]:
            # qsqrt: [L, M, M]
            shape = (num_latent_gps, n_inducing_points, n_inducing_points)
            return tf.linalg.band_part(tf.random.normal(shape, dtype=tf.float64), -1, 0)

        return fn_2
    else:
        raise NotImplementedError


@pytest.fixture(name="whiten", params=[False, True])
def _whiten_fixture(request: SubRequest) -> bool:
    return cast(bool, request.param)


@pytest.fixture(name="num_latent_gps", params=[1, 2])
def _num_latent_gps_fixture(request: SubRequest) -> int:
    return cast(int, request.param)


@pytest.fixture(name="output_dims", params=[1, 5])
def _output_dims_fixture(request: SubRequest) -> int:
    return cast(int, request.param)


ConditionalClosure = Callable[..., tf.Tensor]


@check_shapes(
    "inducing_variable: [M, D, broadcast P]",
    "q_mu: [MxP, R]",
    "q_sqrt: [MxP_or_MxP_N_N...]",
)
def create_conditional(
    *,
    kernel: Kernel,
    inducing_variable: InducingVariables,
    q_mu: TensorType,
    q_sqrt: TensorType,
    whiten: bool,
) -> ConditionalClosure:
    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, R]",
        "return[1]: [batch..., N, R] if (not full_cov) and (not full_output_cov)",
        "return[1]: [batch..., R, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, R, R] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, R, N, R] if full_cov and full_output_cov",
    )
    def conditional_closure(
        Xnew: TensorType, *, full_cov: bool, full_output_cov: bool
    ) -> tf.Tensor:
        return conditional(
            Xnew,
            inducing_variable,
            kernel,
            q_mu,
            q_sqrt=q_sqrt,
            white=whiten,
            full_cov=full_cov,
            full_output_cov=full_output_cov,
        )

    return conditional_closure


def _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
    posterior: AbstractPosterior,
    conditional_closure: ConditionalClosure,
    full_cov: bool,
    full_output_cov: bool,
) -> None:
    Xnew = np.random.randn(13, INPUT_DIMS)

    fused_f_mean, fused_f_cov = posterior.fused_predict_f(
        Xnew, full_cov=full_cov, full_output_cov=full_output_cov
    )

    precomputed_f_mean, precomputed_f_cov = posterior.predict_f(
        Xnew, full_cov=full_cov, full_output_cov=full_output_cov
    )

    conditional_f_mean, conditional_f_cov = conditional_closure(
        Xnew, full_cov=full_cov, full_output_cov=full_output_cov
    )

    np.testing.assert_allclose(fused_f_mean, precomputed_f_mean)
    np.testing.assert_allclose(fused_f_cov, precomputed_f_cov)
    np.testing.assert_array_equal(fused_f_mean, conditional_f_mean)
    np.testing.assert_array_equal(fused_f_cov, conditional_f_cov)


def test_independent_single_output(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    whiten: bool,
    full_cov: bool,
    full_output_cov: bool,
) -> None:
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))

    q_mu = np.random.randn(NUM_INDUCING_POINTS, 1)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, IndependentPosteriorSingleOutput)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_fully_correlated_multi_output(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    output_dims: int,
) -> None:
    """
    The fully correlated posterior has one latent GP.
    """
    kernel = gpflow.kernels.SharedIndependent(
        gpflow.kernels.SquaredExponential(), output_dim=output_dims
    )
    inducing_variable = inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))

    q_mu = np.random.randn(output_dims * NUM_INDUCING_POINTS, 1)
    q_sqrt = q_sqrt_factory(output_dims * NUM_INDUCING_POINTS, 1)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, FullyCorrelatedPosterior)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_independent_multi_output_shk_shi(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    num_latent_gps: int,
    output_dims: int,
) -> None:
    """
    Independent multi-output posterior with a shared kernel and shared inducing points.
    """
    kernel = gpflow.kernels.SharedIndependent(
        gpflow.kernels.SquaredExponential(), output_dim=output_dims
    )
    inducing_variable = gpflow.inducing_variables.SharedIndependentInducingVariables(
        inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, num_latent_gps)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, num_latent_gps)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, IndependentPosteriorMultiOutput)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_independent_multi_output_shk_sei(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    num_latent_gps: int,
    output_dims: int,
) -> None:
    """
    Independent multi-output posterior with a shared kernel and separate inducing points.
    """
    kernel = gpflow.kernels.SharedIndependent(
        gpflow.kernels.SquaredExponential(), output_dim=output_dims
    )
    inducing_variable = gpflow.inducing_variables.SeparateIndependentInducingVariables(
        [
            inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))
            for _ in range(num_latent_gps)
        ]
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, num_latent_gps)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, num_latent_gps)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, IndependentPosteriorMultiOutput)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_independent_multi_output_sek_shi(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    num_latent_gps: int,
    output_dims: int,
) -> None:
    """
    Independent multi-output posterior with separate independent kernels and shared inducing points.
    """
    kernel = gpflow.kernels.SeparateIndependent(
        [gpflow.kernels.SquaredExponential() for _ in range(num_latent_gps)]
    )
    inducing_variable = gpflow.inducing_variables.SharedIndependentInducingVariables(
        inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, num_latent_gps)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, num_latent_gps)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, IndependentPosteriorMultiOutput)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_independent_multi_output_sek_sei(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    num_latent_gps: int,
    output_dims: int,
) -> None:
    """
    Independent multi-output posterior with separate independent kernel and separate inducing points.
    """
    kernel = gpflow.kernels.SeparateIndependent(
        [gpflow.kernels.SquaredExponential() for _ in range(num_latent_gps)]
    )
    inducing_variable = gpflow.inducing_variables.SeparateIndependentInducingVariables(
        [
            inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))
            for _ in range(num_latent_gps)
        ]
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, num_latent_gps)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, num_latent_gps)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, IndependentPosteriorMultiOutput)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_fallback_independent_multi_output_sei(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    output_dims: int,
) -> None:
    """
    Fallback posterior with separate independent inducing variables.

    The FallbackIndependentLatentPosterior is a subclass of the FullyCorrelatedPosterior which
    requires a single latent GP function.
    """
    kernel = gpflow.kernels.LinearCoregionalization(
        [gpflow.kernels.SquaredExponential()], W=tf.random.normal((output_dims, 1))
    )
    inducing_variable = gpflow.inducing_variables.FallbackSeparateIndependentInducingVariables(
        [inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS)) for _ in range(1)]
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, 1)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, FallbackIndependentLatentPosterior)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_fallback_independent_multi_output_shi(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    output_dims: int,
) -> None:
    """
    Fallback posterior with shared independent inducing variables.

    The FallbackIndependentLatentPosterior is a subclass of the FullyCorrelatedPosterior which
    requires a single latent GP function.
    """
    kernel = gpflow.kernels.LinearCoregionalization(
        [gpflow.kernels.SquaredExponential()], W=tf.random.normal((output_dims, 1))
    )
    inducing_variable = gpflow.inducing_variables.FallbackSharedIndependentInducingVariables(
        inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, 1)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, FallbackIndependentLatentPosterior)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_linear_coregionalization_sei(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    num_latent_gps: int,
    output_dims: int,
) -> None:
    """
    Linear coregionalization posterior with separate independent inducing variables.
    """
    kernel = gpflow.kernels.LinearCoregionalization(
        [gpflow.kernels.SquaredExponential() for _ in range(num_latent_gps)],
        W=tf.random.normal((output_dims, num_latent_gps)),
    )
    inducing_variable = gpflow.inducing_variables.SeparateIndependentInducingVariables(
        [
            inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))
            for _ in range(num_latent_gps)
        ]
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, num_latent_gps)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, num_latent_gps)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, LinearCoregionalizationPosterior)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


def test_linear_coregionalization_shi(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    full_cov: bool,
    full_output_cov: bool,
    whiten: bool,
    num_latent_gps: int,
    output_dims: int,
) -> None:
    """
    Linear coregionalization with shared independent inducing variables.
    """
    kernel = gpflow.kernels.LinearCoregionalization(
        [gpflow.kernels.SquaredExponential() for _ in range(num_latent_gps)],
        W=tf.random.normal((output_dims, num_latent_gps)),
    )
    inducing_variable = gpflow.inducing_variables.SharedIndependentInducingVariables(
        inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))
    )

    q_mu = np.random.randn(NUM_INDUCING_POINTS, num_latent_gps)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, num_latent_gps)

    conditional = create_conditional(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    posterior = create_posterior(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
    )
    register_posterior_test(posterior, LinearCoregionalizationPosterior)

    _assert_fused_predict_f_equals_precomputed_predict_f_and_conditional(
        posterior, conditional, full_cov, full_output_cov
    )


@pytest.mark.parametrize(
    "precompute_cache_type", [PrecomputeCacheType.NOCACHE, PrecomputeCacheType.TENSOR]
)
def test_posterior_update_cache_with_variables_no_precompute(
    q_sqrt_factory: QSqrtFactory, whiten: bool, precompute_cache_type: PrecomputeCacheType
) -> None:
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))

    q_mu = np.random.randn(NUM_INDUCING_POINTS, 1)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)

    posterior = IndependentPosteriorSingleOutput(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
        precompute_cache=precompute_cache_type,
    )
    posterior.update_cache(PrecomputeCacheType.VARIABLE)

    assert posterior.cache
    alpha, Qinv = posterior.cache
    assert isinstance(alpha, tf.Variable)
    assert isinstance(Qinv, tf.Variable)


@pytest.mark.parametrize(
    "precompute_cache_type", [PrecomputeCacheType.NOCACHE, PrecomputeCacheType.TENSOR]
)
def test_gpr_posterior_update_cache_with_variables_no_precompute(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    whiten: bool,
    precompute_cache_type: PrecomputeCacheType,
) -> None:
    kernel = gpflow.kernels.SquaredExponential()
    X = np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS)
    Y = np.random.randn(NUM_INDUCING_POINTS, 1)

    posterior = GPRPosterior(
        kernel=kernel,
        data=(X, Y),
        likelihood=Gaussian(0.1),
        precompute_cache=precompute_cache_type,
        mean_function=Zero(),
    )
    posterior.update_cache(PrecomputeCacheType.VARIABLE)
    register_posterior_test(posterior, GPRPosterior)

    assert posterior.cache
    (Kmm_plus_s_inv,) = posterior.cache
    assert isinstance(Kmm_plus_s_inv, tf.Variable)


@pytest.mark.parametrize(
    "precompute_cache_type", [PrecomputeCacheType.NOCACHE, PrecomputeCacheType.TENSOR]
)
def test_sgpr_posterior_update_cache_with_variables_no_precompute(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    whiten: bool,
    precompute_cache_type: PrecomputeCacheType,
) -> None:
    kernel = gpflow.kernels.SquaredExponential()
    X = np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS)
    Y = np.random.randn(NUM_INDUCING_POINTS, 1)
    Z = np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS)

    posterior = SGPRPosterior(
        kernel=kernel,
        data=(X, Y),
        inducing_variable=InducingPoints(Z),
        likelihood=Gaussian(0.1),
        num_latent_gps=1,
        precompute_cache=precompute_cache_type,
        mean_function=Zero(),
    )
    posterior.update_cache(PrecomputeCacheType.VARIABLE)
    register_posterior_test(posterior, SGPRPosterior)

    assert posterior.cache
    alpha, Qinv = posterior.cache
    assert isinstance(alpha, tf.Variable)
    assert isinstance(Qinv, tf.Variable)


@pytest.mark.parametrize(
    "precompute_cache_type", [PrecomputeCacheType.NOCACHE, PrecomputeCacheType.TENSOR]
)
def test_vgp_posterior_update_cache_with_variables_no_precompute(
    register_posterior_test: RegisterPosterior,
    q_sqrt_factory: QSqrtFactory,
    whiten: bool,
    precompute_cache_type: PrecomputeCacheType,
) -> None:
    kernel = gpflow.kernels.SquaredExponential()
    X = np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS)
    q_mu = np.random.randn(NUM_INDUCING_POINTS, 1)
    q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)

    posterior = VGPPosterior(
        kernel=kernel,
        X=X,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        mean_function=Zero(),
        precompute_cache=precompute_cache_type,
    )
    posterior.update_cache(PrecomputeCacheType.VARIABLE)
    register_posterior_test(posterior, VGPPosterior)

    assert posterior.cache
    (Lm,) = posterior.cache
    assert isinstance(Lm, tf.Variable)


def test_posterior_update_cache_with_variables_update_value(
    q_sqrt_factory: QSqrtFactory, whiten: bool
) -> None:
    # setup posterior
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))

    q_mu = tf.Variable(np.random.randn(NUM_INDUCING_POINTS, 1))

    initial_q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)
    if initial_q_sqrt is not None:
        q_sqrt = tf.Variable(initial_q_sqrt)
    else:
        q_sqrt = initial_q_sqrt

    posterior = IndependentPosteriorSingleOutput(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
        precompute_cache=PrecomputeCacheType.TENSOR,
    )
    assert posterior.cache
    initial_alpha, initial_Qinv = posterior.cache

    posterior.update_cache(PrecomputeCacheType.VARIABLE)

    # ensure the values of alpha and Qinv will change
    q_mu.assign_add(tf.ones_like(q_mu))
    if initial_q_sqrt is not None:
        q_sqrt.assign_add(tf.ones_like(q_sqrt))
    posterior.update_cache(PrecomputeCacheType.VARIABLE)

    # assert that the values have changed
    assert posterior.cache
    alpha, Qinv = posterior.cache
    assert not np.allclose(initial_alpha, tf.convert_to_tensor(alpha))
    if initial_q_sqrt is not None:
        assert not np.allclose(initial_Qinv, tf.convert_to_tensor(Qinv))


def test_posterior_update_cache_fails_without_argument(
    q_sqrt_factory: QSqrtFactory, whiten: bool
) -> None:
    # setup posterior
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))

    q_mu = tf.Variable(np.random.randn(NUM_INDUCING_POINTS, 1))

    initial_q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)
    if initial_q_sqrt is not None:
        q_sqrt = tf.Variable(initial_q_sqrt)
    else:
        q_sqrt = initial_q_sqrt

    posterior = IndependentPosteriorSingleOutput(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
        precompute_cache=None,
    )
    assert posterior.cache is None

    with pytest.raises(ValueError):
        posterior.update_cache()

    posterior.update_cache(PrecomputeCacheType.TENSOR)
    assert posterior.cache
    assert all(isinstance(c, tf.Tensor) for c in posterior.cache)

    posterior.update_cache(PrecomputeCacheType.NOCACHE)
    assert posterior._precompute_cache == PrecomputeCacheType.NOCACHE
    assert posterior.cache is None

    posterior.update_cache(PrecomputeCacheType.TENSOR)  # set posterior._precompute_cache
    assert posterior._precompute_cache == PrecomputeCacheType.TENSOR

    posterior.cache = None  # clear again
    posterior.update_cache()  # does not raise an exception
    assert posterior.cache
    assert all(isinstance(c, tf.Tensor) for c in posterior.cache)


def test_posterior_create_with_variables_update_cache_works(
    q_sqrt_factory: QSqrtFactory, whiten: bool
) -> None:
    # setup posterior
    kernel = gpflow.kernels.SquaredExponential()
    inducing_variable = inducingpoint_wrapper(np.random.randn(NUM_INDUCING_POINTS, INPUT_DIMS))

    q_mu = tf.Variable(np.random.randn(NUM_INDUCING_POINTS, 1))

    initial_q_sqrt = q_sqrt_factory(NUM_INDUCING_POINTS, 1)
    if initial_q_sqrt is not None:
        q_sqrt = tf.Variable(initial_q_sqrt)
    else:
        q_sqrt = initial_q_sqrt

    posterior = IndependentPosteriorSingleOutput(
        kernel=kernel,
        inducing_variable=inducing_variable,
        q_mu=q_mu,
        q_sqrt=q_sqrt,
        whiten=whiten,
        precompute_cache=PrecomputeCacheType.VARIABLE,
    )
    assert posterior.cache
    assert all(isinstance(c, tf.Variable) for c in posterior.cache)

    cache = posterior.cache

    posterior.update_cache()

    assert posterior.cache is cache


# tests/gpflow/quadrature/ndiagquad_old.py

from typing import Callable, Iterable, List, Sequence, Tuple, Union, cast

import numpy as np
import tensorflow as tf

from gpflow.base import TensorType
from gpflow.experimental.check_shapes import check_shapes
from gpflow.quadrature.deprecated import mvhermgauss


@check_shapes(
    "Fmu: [broadcast Din, N...]",
    "Fvar: [broadcast Din, N...]",
    "Ys.values(): [N...]",
    "return: [broadcast Dout, N...]",
)
def ndiagquad(
    funcs: Union[Callable[..., tf.Tensor], Iterable[Callable[..., tf.Tensor]]],
    H: int,
    Fmu: Union[TensorType, List[TensorType], Tuple[TensorType, ...]],
    Fvar: Union[TensorType, List[TensorType], Tuple[TensorType, ...]],
    logspace: bool = False,
    **Ys: TensorType,
) -> tf.Tensor:
    """
    Computes N Gaussian expectation integrals of one or more functions
    using Gauss-Hermite quadrature. The Gaussians must be independent.
    The means and variances of the Gaussians are specified by Fmu and Fvar.
    The N-integrals are assumed to be taken wrt the last dimensions of Fmu, Fvar.
    :param funcs: the integrand(s):
        Callable or Iterable of Callables that operates elementwise
    :param H: number of Gauss-Hermite quadrature points
    :param Fmu: array/tensor or `Din`-tuple/list thereof
    :param Fvar: array/tensor or `Din`-tuple/list thereof
    :param logspace: if True, funcs are the log-integrands and this calculates
        the log-expectation of exp(funcs)
    :param **Ys: arrays/tensors; deterministic arguments to be passed by name
    Fmu, Fvar, Ys should all have same shape, with overall size `N`
    :return: shape is the same as that of the first Fmu
    """
    if isinstance(Fmu, (tuple, list)):
        assert isinstance(Fvar, (tuple, list))  # Hint for mypy.
        Din = len(Fmu)

        def unify(f_list: Sequence[TensorType]) -> tf.Tensor:
            """Stack a list of means/vars into a full block."""
            return tf.reshape(
                tensor=tf.concat([tf.reshape(f, shape=(-1, 1)) for f in f_list], axis=1),
                shape=(-1, 1, Din),
            )

        shape = tf.shape(Fmu[0])
        Fmu, Fvar = map(unify, [Fmu, Fvar])  # both [N, 1, Din]

        print(Fmu)
        print(Fvar)
    else:
        Din = 1
        shape = tf.shape(Fmu)
        Fmu, Fvar = [tf.reshape(f, (-1, 1, 1)) for f in [Fmu, Fvar]]

    Fmu = cast(TensorType, Fmu)
    Fvar = cast(TensorType, Fvar)

    xn, wn = mvhermgauss(H, Din)
    # xn: H**Din x Din, wn: H**Din

    gh_x = xn.reshape(1, -1, Din)  # [1, H]**Din x Din
    Xall = gh_x * tf.sqrt(2.0 * Fvar) + Fmu  # [N, H]**Din x Din
    Xs = [Xall[:, :, i] for i in range(Din)]  # [N, H]**Din  each

    gh_w = wn * np.pi ** (-0.5 * Din)  # H**Din x 1

    for name, Y in Ys.items():
        Y = tf.reshape(Y, (-1, 1))
        Y = tf.tile(Y, [1, H ** Din])  # broadcast Y to match X
        # without the tiling, some calls such as tf.where() (in bernoulli) fail
        Ys[name] = Y  # now [N, H]**Din

    def eval_func(f: Callable[..., tf.Tensor]) -> tf.Tensor:
        feval = f(*Xs, **Ys)  # f should be elementwise: return shape [N, H]**Din
        if logspace:
            log_gh_w = np.log(gh_w.reshape(1, -1))
            result = tf.reduce_logsumexp(feval + log_gh_w, axis=1)
        else:
            result = tf.linalg.matmul(feval, gh_w.reshape(-1, 1))
        return tf.reshape(result, shape)

    if isinstance(funcs, Iterable):
        return [eval_func(f) for f in funcs]

    return eval_func(funcs)


# tests/gpflow/quadrature/test_quadrature.py

# Copyright 2018 the GPflow authors.
#
# 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 typing import cast

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose

import gpflow.quadrature as quadrature
from gpflow.base import AnyNDArray, TensorType


@pytest.mark.parametrize("mu", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var", [np.array([3.0, 3.5])])
def test_diagquad_1d(mu: TensorType, var: TensorType) -> None:
    num_gauss_hermite_points = 25
    quad = quadrature.ndiagquad([lambda *X: tf.exp(X[0])], num_gauss_hermite_points, [mu], [var])
    expected = np.exp(mu + var / 2)
    assert_allclose(quad[0], expected)


@pytest.mark.parametrize("mu1", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var1", [np.array([3.0, 3.5])])
@pytest.mark.parametrize("mu2", [np.array([-2.0, 0.3])])
@pytest.mark.parametrize("var2", [np.array([4.0, 4.2])])
def test_diagquad_2d(mu1: TensorType, var1: TensorType, mu2: TensorType, var2: TensorType) -> None:
    alpha = 2.5
    # using logspace=True we can reduce this, see test_diagquad_logspace
    num_gauss_hermite_points = 35
    quad = quadrature.ndiagquad(
        lambda *X: tf.exp(X[0] + alpha * X[1]),
        num_gauss_hermite_points,
        [mu1, mu2],
        [var1, var2],
    )
    expected = np.exp(mu1 + var1 / 2 + alpha * mu2 + alpha ** 2 * var2 / 2)
    assert_allclose(quad, expected)


@pytest.mark.parametrize("mu1", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var1", [np.array([3.0, 3.5])])
@pytest.mark.parametrize("mu2", [np.array([-2.0, 0.3])])
@pytest.mark.parametrize("var2", [np.array([4.0, 4.2])])
def test_diagquad_logspace(
    mu1: TensorType, var1: TensorType, mu2: TensorType, var2: TensorType
) -> None:
    alpha = 2.5
    num_gauss_hermite_points = 25
    quad = quadrature.ndiagquad(
        lambda *X: (X[0] + alpha * X[1]),
        num_gauss_hermite_points,
        [mu1, mu2],
        [var1, var2],
        logspace=True,
    )
    expected = mu1 + var1 / 2 + alpha * mu2 + alpha ** 2 * var2 / 2
    assert_allclose(quad, expected)


@pytest.mark.parametrize("mu1", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var1", [np.array([3.0, 3.5])])
def test_diagquad_with_kwarg(mu1: TensorType, var1: TensorType) -> None:
    alpha: AnyNDArray = np.array([2.5, -1.3])
    num_gauss_hermite_points = 25
    quad = quadrature.ndiagquad(
        lambda X, Y: tf.exp(X * Y), num_gauss_hermite_points, mu1, var1, Y=alpha
    )
    expected = np.exp(alpha * mu1 + alpha ** 2 * var1 / 2)
    assert_allclose(quad, expected)


def test_ndiagquad_does_not_throw_error() -> None:
    """
    Check that the autograph=False for quadrature.ndiagquad does not throw an error.
    Regression test for https://github.com/GPflow/GPflow/issues/1547.
    """

    @tf.function(autograph=False)
    def func_ndiagquad_autograph_false() -> tf.Tensor:
        mu: AnyNDArray = np.array([1.0, 1.3])
        var: AnyNDArray = np.array([3.0, 3.5])
        num_gauss_hermite_points = 25
        return quadrature.ndiagquad(
            [lambda *X: tf.exp(X[0])], num_gauss_hermite_points, [mu], [var]
        )

    func_ndiagquad_autograph_false()


def test_quadrature_autograph() -> None:
    """
    Check that the return value is equal with and without Autograph
    Regression test for https://github.com/GPflow/GPflow/issues/1547.
    """

    def compute(autograph: bool) -> AnyNDArray:
        @tf.function(autograph=autograph)
        def func() -> tf.Tensor:
            mu: AnyNDArray = np.array([1.0, 1.3])
            var: AnyNDArray = np.array([3.0, 3.5])
            num_gauss_hermite_points = 25
            return quadrature.ndiagquad(
                [lambda *X: tf.exp(X[0])], num_gauss_hermite_points, [mu], [var]
            )

        (result,) = func()
        return cast(AnyNDArray, result.numpy())

    np.testing.assert_equal(
        compute(autograph=True),
        compute(autograph=False),
    )


# tests/gpflow/quadrature/test_quadrature_equivalence.py

# Copyright 2018 the GPflow authors.
#
# 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.

import numpy as np
import pytest
import tensorflow as tf
from ndiagquad_old import ndiagquad as ndiagquad_old
from numpy.testing import assert_allclose

from gpflow.base import AnyNDArray, TensorType
from gpflow.quadrature import ndiagquad


@pytest.mark.parametrize("mu", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var", [np.array([3.0, 3.5])])
def test_diagquad_1d(mu: TensorType, var: TensorType) -> None:
    num_gauss_hermite_points = 25
    quad = ndiagquad([lambda *X: tf.exp(X[0])], num_gauss_hermite_points, [mu], [var])
    quad_old = ndiagquad_old([lambda *X: tf.exp(X[0])], num_gauss_hermite_points, [mu], [var])
    assert_allclose(quad[0], quad_old[0])


@pytest.mark.parametrize("mu1", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var1", [np.array([3.0, 3.5])])
@pytest.mark.parametrize("mu2", [np.array([-2.0, 0.3])])
@pytest.mark.parametrize("var2", [np.array([4.0, 4.2])])
def test_diagquad_2d(mu1: TensorType, var1: TensorType, mu2: TensorType, var2: TensorType) -> None:
    alpha = 2.5
    # using logspace=True we can reduce this, see test_diagquad_logspace
    num_gauss_hermite_points = 35
    quad = ndiagquad(
        lambda *X: tf.exp(X[0] + alpha * X[1]),
        num_gauss_hermite_points,
        [mu1, mu2],
        [var1, var2],
    )
    quad_old = ndiagquad_old(
        lambda *X: tf.exp(X[0] + alpha * X[1]),
        num_gauss_hermite_points,
        [mu1, mu2],
        [var1, var2],
    )
    assert_allclose(quad, quad_old)


@pytest.mark.parametrize("mu1", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var1", [np.array([3.0, 3.5])])
@pytest.mark.parametrize("mu2", [np.array([-2.0, 0.3])])
@pytest.mark.parametrize("var2", [np.array([4.0, 4.2])])
def test_diagquad_logspace(
    mu1: TensorType, var1: TensorType, mu2: TensorType, var2: TensorType
) -> None:
    alpha = 2.5
    num_gauss_hermite_points = 25
    quad = ndiagquad(
        lambda *X: (X[0] + alpha * X[1]),
        num_gauss_hermite_points,
        [mu1, mu2],
        [var1, var2],
        logspace=True,
    )
    quad_old = ndiagquad_old(
        lambda *X: (X[0] + alpha * X[1]),
        num_gauss_hermite_points,
        [mu1, mu2],
        [var1, var2],
        logspace=True,
    )
    assert_allclose(quad, quad_old)


@pytest.mark.parametrize("mu1", [np.array([1.0, 1.3])])
@pytest.mark.parametrize("var1", [np.array([3.0, 3.5])])
def test_diagquad_with_kwarg(mu1: TensorType, var1: TensorType) -> None:
    alpha: AnyNDArray = np.array([2.5, -1.3])
    num_gauss_hermite_points = 25
    quad = ndiagquad(lambda X, Y: tf.exp(X * Y), num_gauss_hermite_points, mu1, var1, Y=alpha)
    quad_old = ndiagquad_old(
        lambda X, Y: tf.exp(X * Y), num_gauss_hermite_points, mu1, var1, Y=alpha
    )
    assert_allclose(quad, quad_old)


# tests/gpflow/utilities/test_bijectors.py

from typing import Optional, Type

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp

from gpflow.config import Config, as_context
from gpflow.utilities import positive, triangular, triangular_size


@pytest.mark.parametrize(
    "env_lower, override_lower",
    [
        (0.1, None),  # ensure default from config is applied
        (0.0, 0.2),  # ensure override is applied
        (0.3, 0.4),  # ensure local overrides config
    ],
)
def test_positive_lower(env_lower: float, override_lower: Optional[float]) -> None:
    expected_lower = override_lower or env_lower
    with as_context(Config(positive_bijector="softplus", positive_minimum=env_lower)):
        bijector = positive(lower=override_lower)
        assert isinstance(bijector, tfp.bijectors.Chain)
        assert np.isclose(bijector.bijectors[0].shift, expected_lower)


@pytest.mark.parametrize(
    "env_bijector, override_bijector, expected_class",
    [
        ("softplus", None, tfp.bijectors.Softplus),
        ("softplus", "Exp", tfp.bijectors.Exp),
        ("exp", None, tfp.bijectors.Exp),
        ("exp", "Softplus", tfp.bijectors.Softplus),
    ],
)
def test_positive_bijector(
    env_bijector: str,
    override_bijector: Optional[str],
    expected_class: Type[tfp.bijectors.Bijector],
) -> None:
    with as_context(Config(positive_bijector=env_bijector, positive_minimum=0.0)):
        bijector = positive(base=override_bijector)
        assert isinstance(bijector, expected_class)


def test_positive_calculation_order() -> None:
    value, lower = -10.0, 10.0
    expected = np.exp(value) + lower
    with as_context(Config(positive_bijector="exp", positive_minimum=lower)):
        result = positive()(value).numpy()
    assert np.isclose(result, expected)
    assert result >= lower


def test_triangular() -> None:
    assert isinstance(triangular(), tfp.bijectors.FillTriangular)


@pytest.mark.parametrize(
    "n,expected",
    [
        (0, 0),
        (1, 1),
        (2, 3),
        (3, 6),
        (4, 10),
    ],
)
def test_triangular_size(n: int, expected: int) -> None:
    actual = triangular_size(tf.constant(n))
    assert actual.dtype.is_integer
    assert expected == actual


# tests/gpflow/utilities/test_deepcopy.py

import copy
import pickle
from typing import Any, Type

import pytest
import tensorflow as tf
import tensorflow_probability as tfp
from packaging.version import Version

import gpflow
from gpflow.base import TensorType
from gpflow.experimental.check_shapes import check_shapes


class A(tf.Module):
    def __init__(self) -> None:
        self.var = tf.Variable([1.0])
        shift = tf.Variable(1e-6)
        self.bijector = tfp.bijectors.Chain([tfp.bijectors.Softplus(), tfp.bijectors.Shift(shift)])

    @check_shapes(
        "x: [any...]",
        "return: [any...]",
    )
    def __call__(self, x: TensorType) -> tf.Tensor:
        return self.bijector(x)


class B(tf.Module):
    def __init__(self) -> None:
        self.var = tf.Variable([2.0])
        self.a = A()

    @check_shapes(
        "x: [any...]",
        "return: [any...]",
    )
    def __call__(self, x: TensorType) -> tf.Tensor:
        return self.a(x)


class C(tf.Module):
    def __init__(self) -> None:
        self.var = gpflow.Parameter([2.0], transform=gpflow.utilities.positive(lower=1e-6))


class D(tf.Module):
    def __init__(self) -> None:
        self.var = gpflow.Parameter([10.0], transform=gpflow.utilities.positive())
        self.var2 = gpflow.Parameter([5.0])
        self.c = C()


class NestedModule(tf.Module):
    def __init__(self, module: tf.Module) -> None:
        self.module = module


@pytest.mark.parametrize("module", [A(), B()])
def test_clears_bijector_cache_and_deepcopy(module: Type[Any]) -> None:
    """
    With each forward pass through a bijector, a cache is stored inside which prohibits the deepcopy of the bijector.
    This is due to the fact that HashableWeakRef objects are not pickle-able, which raises a TypeError. Alternatively,
    one can make use of `deepcopy_component` to deepcopy a module containing used bijectors.
    """
    input = 1.0
    _ = module(input)

    if Version(tfp.__version__) >= Version("0.12"):
        module_copy = copy.deepcopy(module)
    else:
        # older versions of tensorflow_probability required us to maintain a workaround
        with pytest.raises(TypeError):
            copy.deepcopy(module)
        module_copy = gpflow.utilities.deepcopy(module)

    assert module.var == module_copy.var
    assert module.var is not module_copy.var
    module_copy.var.assign([5.0])
    assert module.var != module_copy.var


def test_freeze() -> None:
    module = NestedModule(NestedModule(A()))
    module_frozen = gpflow.utilities.freeze(module)
    assert len(module.variables) == 2
    assert module_frozen.variables == ()
    assert isinstance(module.module.module.var, tf.Variable)
    assert isinstance(module_frozen.module.module.var, tf.Tensor)


def test_pickle_frozen() -> None:
    """
    Regression test for the bug described in GPflow/GPflow#1338
    """
    module = D()
    module_frozen = gpflow.utilities.freeze(module)

    pickled = pickle.dumps(module_frozen)
    loaded = pickle.loads(pickled)

    assert loaded.var == module_frozen.var
    assert loaded.var2 == module_frozen.var2
    assert loaded.c.var == module_frozen.c.var


# tests/gpflow/utilities/test_misc.py

# Copyright 2017-2021 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 typing import Any

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp

from gpflow.base import Parameter
from gpflow.utilities import is_variable


@pytest.mark.parametrize(
    "value,expected",
    [
        ([[0.1, 0.2], [0.3, 0.4]], False),
        (np.array(5), False),
        (tf.constant(5), False),
        (tf.Variable(5), True),
        (tfp.util.TransformedVariable(5, tfp.bijectors.Identity()), True),
        (Parameter(5), True),
    ],
    ids=lambda x: x if isinstance(x, bool) else type(x).__name__,
)
def test_is_variable(value: Any, expected: bool) -> None:
    assert expected == is_variable(value)


# tests/gpflow/utilities/test_model_utils.py

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import TensorType
from gpflow.utilities import add_likelihood_noise_cov, add_noise_cov, assert_params_false


def test_assert_params_false__False() -> None:
    assert_params_false(test_assert_params_false__False, foo=False, bar=False)


def test_assert_params_false__True() -> None:
    with pytest.raises(NotImplementedError):
        assert_params_false(test_assert_params_false__True, foo=False, bar=True)


@pytest.mark.parametrize(
    "input_tensor,variance,expected_tensor",
    [
        (
            tf.constant([[1.0, 0.5], [0.5, 1.0]]),
            gpflow.Parameter(1.0, dtype=tf.float32),
            tf.constant([[2.0, 0.5], [0.5, 2.0]]),
        ),
        (
            tf.constant(
                [
                    [
                        [
                            [0.0, 0.0],
                            [0.0, 0.0],
                        ],
                    ],
                    [
                        [
                            [1.0, 1.0],
                            [1.0, 1.0],
                        ],
                    ],
                ]
            ),
            gpflow.Parameter(0.5, dtype=tf.float32),
            tf.constant(
                [
                    [
                        [
                            [0.5, 0.0],
                            [0.0, 0.5],
                        ],
                    ],
                    [
                        [
                            [1.5, 1.0],
                            [1.0, 1.5],
                        ],
                    ],
                ]
            ),
        ),
        (
            tf.constant([[1.0, 0.5], [0.5, 1.0]]),
            gpflow.Parameter([1.0, 2.0], dtype=tf.float32),
            tf.constant([[2.0, 0.5], [0.5, 3.0]]),
        ),
    ],
)
def test_add_noise_cov(
    input_tensor: TensorType, variance: TensorType, expected_tensor: TensorType
) -> None:
    actual_tensor = add_noise_cov(input_tensor, variance)
    np.testing.assert_allclose(actual_tensor, expected_tensor)


def test_add_likelihood_noise_cov() -> None:
    K = tf.constant(
        [
            [0.11, 0.12, 0.13],
            [0.21, 0.22, 0.23],
            [0.31, 0.32, 0.33],
        ],
        dtype=gpflow.default_float(),
    )
    variance = gpflow.functions.Linear(A=[[0.2]], b=0.1)
    likelihood = gpflow.likelihoods.Gaussian(variance=variance)
    X = tf.constant([[2.0], [4.0], [3.0]], dtype=gpflow.default_float())
    actual_tensor = add_likelihood_noise_cov(K, likelihood, X)

    np.testing.assert_allclose(
        [
            [0.61, 0.12, 0.13],
            [0.21, 1.12, 0.23],
            [0.31, 0.32, 1.03],
        ],
        actual_tensor,
    )


# tests/gpflow/utilities/test_multipledispatch.py

import re
import warnings
from typing import Any, Type

import multipledispatch
import pytest
import tensorflow as tf
from _pytest.capture import CaptureFixture
from packaging.version import Version

import gpflow
from gpflow.utilities import Dispatcher


class A1:
    pass


class A2(A1):
    pass


class B1:
    pass


class B2(B1):
    pass


def get_test_fn() -> Dispatcher:
    test_fn = Dispatcher("test_fn")

    @test_fn.register(A1, B1)
    def test_a1_b1(x: A1, y: B1) -> str:
        return "a1-b1"

    @test_fn.register(A2, B1)
    def test_a2_b1(x: A2, y: B1) -> str:
        return "a2-b1"

    @test_fn.register(A1, B2)
    def test_a1_b2(x: A1, y: B2) -> str:
        return "a1-b2"

    return test_fn


def test_our_multipledispatch() -> None:
    test_fn = get_test_fn()

    assert test_fn(A1(), B1()) == "a1-b1"
    assert test_fn(A2(), B1()) == "a2-b1"
    assert test_fn(A1(), B2()) == "a1-b2"

    # test the ambiguous case:

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")

        assert test_fn(A2(), B2()) == "a1-b2"  # the last definition wins

        assert len(w) == 1
        assert issubclass(w[0].category, multipledispatch.conflict.AmbiguityWarning)

    # test that adding the child-child definition removes ambiguity warning:

    @test_fn.register(A2, B2)
    def test_a2_b2(x: A2, y: B2) -> str:
        return "a2-b2"

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")

        assert test_fn(A2(), B2()) == "a2-b2"

        assert len(w) == 0


def test_dispatcher__no_match() -> None:
    test_fn = get_test_fn()

    with pytest.raises(NotImplementedError):
        test_fn(3, "foo")

    assert None is test_fn.dispatch(int, str)

    with pytest.raises(NotImplementedError):
        test_fn.dispatch_or_raise(int, str)


@pytest.mark.parametrize(
    "Dispatcher, expect_autograph_warning",
    [
        (
            multipledispatch.Dispatcher,
            True,
        ),  # checks test is written correctly and captures warnings generated by base Dispatcher
        (gpflow.utilities.Dispatcher, False),  # no warnings with our custom Dispatcher
    ],
)
def test_dispatcher_autograph_warnings(
    capsys: CaptureFixture[str], Dispatcher: Type[Any], expect_autograph_warning: bool
) -> None:
    if Dispatcher is multipledispatch.Dispatcher and (Version(tf.__version__) >= Version("2.3.0")):
        pytest.skip("In TensorFlow >= 2.3, multipledispatch.Dispatcher no longer works at all")

    tf.autograph.set_verbosity(0, alsologtostdout=True)  # to be able to capture it using capsys

    test_fn = Dispatcher("test_fn")

    # generator would only be invoked when defining for base class...
    @test_fn.register(gpflow.inducing_variables.InducingVariables)
    def test_iv(x: gpflow.inducing_variables.InducingVariables) -> tf.Tensor:
        return tf.reduce_sum(x.Z)

    test_fn_compiled = tf.function(test_fn)  # with autograph=True by default

    # ...but calling using subclass
    result = test_fn_compiled(gpflow.inducing_variables.InducingPoints([[1.0, 2.0]]))
    assert result.numpy() == 3.0  # expect computation to work either way

    captured = capsys.readouterr()

    tf_warning = "WARNING:.*Entity .* appears to be a generator function. It will not be converted by AutoGraph."
    assert bool(re.match(tf_warning, captured.out)) == expect_autograph_warning


# tests/gpflow/utilities/test_ops.py

from typing import Union, cast

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import check_shapes
from gpflow.utilities.ops import difference_matrix


@check_shapes(
    "X: [N, D]",
    "Q: []",
    "return: [N, Q]",
)
def pca_reduce(X: Union[AnyNDArray, tf.Tensor], Q: int) -> AnyNDArray:
    """
    A helpful function for linearly reducing the dimensionality of the data X
    to Q.
    :param X: data array of size N (number of points) x D (dimensions)
    :param Q: Number of latent dimensions, Q < D
    :return: PCA projection array of size N x Q.
    """
    if Q > X.shape[1]:  # pragma: no cover
        raise ValueError("Cannot have more latent dimensions than observed")
    if isinstance(X, tf.Tensor):
        X = X.numpy()
    assert isinstance(X, np.ndarray)  # Hint for numpy
    # TODO why not use tf.linalg.eigh?
    evals, evecs = np.linalg.eigh(np.cov(X.T))
    W = evecs[:, -Q:]
    return cast(AnyNDArray, (X - X.mean(0)).dot(W))


@pytest.mark.parametrize("N", [3, 7])
@pytest.mark.parametrize("D", [2, 5, 9])
@pytest.mark.parametrize("Q", [2, 5, 9])
def test_pca_reduce_numpy_equivalence(N: int, D: int, Q: int) -> None:
    X = np.random.randn(N, D)

    if Q > D:
        with pytest.raises(ValueError):
            gpflow.utilities.ops.pca_reduce(tf.convert_to_tensor(X), Q)

    else:
        np_result = pca_reduce(X, Q)
        tf_result = gpflow.utilities.ops.pca_reduce(tf.convert_to_tensor(X), Q).numpy()
        assert np_result.shape == tf_result.shape == (N, Q)

        for i in range(Q):
            # PCA does not necessarily preserve the overall sign, so also accept it to flip
            tf_column = tf_result[:, i]
            np_column = np_result[:, i]
            assert np.allclose(tf_column, np_column) or np.allclose(tf_column, -np_column)


def test_difference_matrix_broadcasting_symmetric() -> None:
    X = np.random.randn(5, 4, 3, 2)
    d = difference_matrix(X, None)
    assert d.shape == (5, 4, 3, 3, 2)


def test_difference_matrix_broadcasting_cross() -> None:
    X = np.random.randn(2, 3, 4, 5)
    X2 = np.random.randn(8, 7, 6, 5)
    d = difference_matrix(X, X2)
    assert d.shape == (2, 3, 4, 8, 7, 6, 5)


# tests/gpflow/utilities/test_parameter_or_function.py

# Copyright 2017 the GPflow authors.
#
# 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.

import numpy as np
import pytest
import tensorflow_probability as tfp
from numpy.testing import assert_allclose

from gpflow.base import Parameter
from gpflow.functions import Linear
from gpflow.utilities.parameter_or_function import (
    evaluate_parameter_or_function,
    prepare_parameter_or_function,
)


def test_prepare_parameter_or_function__constant__no_bound() -> None:
    initial = 5.0
    param = prepare_parameter_or_function(initial)
    assert isinstance(param, Parameter)
    assert isinstance(param.transform, tfp.bijectors.Identity)
    assert_allclose(initial, param)

    X = np.array(
        [
            [0.0, 0.0],
            [0.0, 1.0],
            [0.0, 2.0],
            [1.0, 0.0],
            [1.0, 1.0],
            [1.0, 2.0],
        ]
    )
    assert_allclose(initial, evaluate_parameter_or_function(param, X))


def test_prepare_parameter_or_function__constant__bound() -> None:
    initial = 5.0

    with pytest.raises(Exception):
        prepare_parameter_or_function(initial, lower_bound=initial + 1e-3)

    lower_bound = initial - 1e-3
    param = prepare_parameter_or_function(initial, lower_bound=lower_bound)
    assert isinstance(param, Parameter)
    assert isinstance(param.transform, tfp.bijectors.Chain)
    assert_allclose(initial, param)

    X = np.array(
        [
            [0.0, 0.0],
            [0.0, 1.0],
            [0.0, 2.0],
            [1.0, 0.0],
            [1.0, 1.0],
            [1.0, 2.0],
        ]
    )
    assert_allclose(initial, evaluate_parameter_or_function(param, X, lower_bound=lower_bound))


def test_prepare_parameter_or_function__function__no_bound() -> None:
    initial = Linear([[0.5], [2.0]], 1.0)
    func = prepare_parameter_or_function(initial)
    assert initial is func

    X = np.array(
        [
            [0.0, 0.0],
            [0.0, 1.0],
            [0.0, 2.0],
            [1.0, 0.0],
            [1.0, 1.0],
            [1.0, 2.0],
        ]
    )
    assert_allclose(
        [[1.0], [3.0], [5.0], [1.5], [3.5], [5.5]], evaluate_parameter_or_function(func, X)
    )


def test_prepare_parameter_or_function__function__bound() -> None:
    initial = Linear([[0.5], [2.0]], 1.0)
    lower_bound = 3.2
    func = prepare_parameter_or_function(initial, lower_bound=lower_bound)
    assert initial is func

    X = np.array(
        [
            [0.0, 0.0],
            [0.0, 1.0],
            [0.0, 2.0],
            [1.0, 0.0],
            [1.0, 1.0],
            [1.0, 2.0],
        ]
    )
    assert_allclose(
        [[3.2], [3.2], [5.0], [3.2], [3.5], [5.5]],
        evaluate_parameter_or_function(func, X, lower_bound=lower_bound),
    )


# tests/gpflow/utilities/test_set_trainable.py

import tensorflow as tf

from gpflow import Parameter, set_trainable


def _module() -> tf.Module:
    class _Mod(tf.Module):
        def __init__(self) -> None:
            super().__init__()
            self.var = tf.Variable(0.0)
            self.param = Parameter(0.0)

    module = _Mod()

    assert len(module.trainable_variables) == 2
    assert len(module.variables) == 2

    return module


def test_can_set_not_trainable() -> None:
    module = _module()
    set_trainable(module, False)
    assert len(module.trainable_variables) == 0


def test_can_set_not_trainable_then_trainable_again() -> None:
    module = _module()
    set_trainable(module, False)
    set_trainable(module, True)
    assert len(module.trainable_variables) == len(module.variables)


def test_can_set_not_trainable_iterable() -> None:
    modules = [_module(), _module(), _module()]
    set_trainable(modules, False)
    assert all(len(m.trainable_variables) == 0 for m in modules)


def test_can_set_not_trainable_then_trainable_iterable() -> None:
    modules = [_module(), _module(), _module()]
    set_trainable(modules, False)
    set_trainable(modules, True)
    assert all(len(m.trainable_variables) == len(m.variables) for m in modules)


# tests/gpflow/utilities/test_training_loop.py

from typing import Any

import numpy as np
import tensorflow as tf

import gpflow
from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.utilities import training_loop


class Datum:
    rng = np.random.RandomState(123)
    cs = ShapeChecker().check_shape

    N = 13
    X = cs(np.linspace(0, 10, N)[:, None], "[N, 1]")
    noise_scale = 0.01
    Y = cs(np.sin(X) + rng.randn(N, 1) * noise_scale, "[N, 1]")
    data = (X, Y)


def create_model() -> gpflow.models.GPModel:
    return gpflow.models.GPR(
        Datum.data, gpflow.kernels.SquaredExponential(), noise_variance=Datum.noise_scale ** 2
    )


def assert_models_close(
    m: gpflow.models.GPModel, mref: gpflow.models.GPModel, **tol_kwargs: Any
) -> None:
    np.testing.assert_allclose(
        m.kernel.variance.numpy(), mref.kernel.variance.numpy(), **tol_kwargs
    )
    np.testing.assert_allclose(
        m.kernel.lengthscales.numpy(), mref.kernel.lengthscales.numpy(), **tol_kwargs
    )
    np.testing.assert_allclose(
        m.likelihood.variance.numpy(), mref.likelihood.variance.numpy(), **tol_kwargs
    )


def test_training_loop_compiles() -> None:
    m1 = create_model()
    m2 = create_model()
    training_loop(
        m1.training_loss, tf.optimizers.Adam(), m1.trainable_variables, maxiter=50, compile=True
    )
    training_loop(
        m2.training_loss, tf.optimizers.Adam(), m2.trainable_variables, maxiter=50, compile=False
    )
    assert_models_close(m1, m2)


def test_training_loop_converges() -> None:
    m = create_model()
    mref = create_model()
    gpflow.optimizers.Scipy().minimize(mref.training_loss, mref.trainable_variables)
    training_loop(
        m.training_loss,
        tf.optimizers.Adam(learning_rate=0.01),
        m.trainable_variables,
        maxiter=5000,
        compile=True,
    )
    assert_models_close(m, mref, rtol=1e-5)


# tests/gpflow/utilities/test_traversal.py

# Copyright 2018 the GPflow authors.
#
# 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 typing import Any, Callable, Mapping, Optional, Type

import numpy as np
import pytest
import tensorflow as tf
from _pytest.fixtures import SubRequest
from packaging.version import Version

import gpflow
from gpflow.base import AnyNDArray
from gpflow.config import Config, as_context
from gpflow.utilities import set_trainable
from gpflow.utilities.traversal import (
    _merge_leaf_components,
    leaf_components,
    tabulate_module_summary,
)

rng = np.random.RandomState(0)


class Data:
    H0 = 5
    H1 = 2
    M = 10
    D = 1
    Z: AnyNDArray = 0.5 * np.ones((M, 1))
    ls = 2.0
    var = 1.0


# ------------------------------------------
# Helpers
# ------------------------------------------


class A(tf.Module):
    def __init__(self, name: Optional[str] = None) -> None:
        super().__init__(name)
        self.var_trainable = tf.Variable(tf.zeros((2, 2, 1)), trainable=True)
        self.var_fixed = tf.Variable(tf.ones((2, 2, 1)), trainable=False)


class B(tf.Module):
    def __init__(self, name: Optional[str] = None) -> None:
        super().__init__(name)
        self.submodule_list = [A(), A()]
        self.submodule_dict = dict(a=A(), b=A())
        self.var_trainable = tf.Variable(tf.zeros((2, 2, 1)), trainable=True)
        self.var_fixed = tf.Variable(tf.ones((2, 2, 1)), trainable=False)


class C(tf.keras.Model):
    def __init__(self, name: Optional[str] = None) -> None:
        super().__init__(name)
        self.variable = tf.Variable(tf.zeros((2, 2, 1)), trainable=True)
        self.param = gpflow.Parameter(0.0)
        self.dense = tf.keras.layers.Dense(5)


def create_kernel() -> gpflow.kernels.Kernel:
    kern = gpflow.kernels.SquaredExponential(lengthscales=Data.ls, variance=Data.var)
    set_trainable(kern.lengthscales, False)
    return kern


def create_compose_kernel() -> gpflow.kernels.Kernel:
    kernel = gpflow.kernels.Product(
        [
            gpflow.kernels.Sum([create_kernel(), create_kernel()]),
            gpflow.kernels.Sum([create_kernel(), create_kernel()]),
        ]
    )
    return kernel


def create_model() -> gpflow.models.GPModel:
    kernel = create_kernel()
    model = gpflow.models.SVGP(
        kernel=kernel,
        likelihood=gpflow.likelihoods.Gaussian(variance_lower_bound=0.0),
        inducing_variable=Data.Z,
        q_diag=True,
    )
    set_trainable(model.q_mu, False)
    return model


# ------------------------------------------
# Reference
# ------------------------------------------

example_tf_module_variable_dict = {
    "A.var_trainable": {
        "value": np.zeros((2, 2, 1)),
        "trainable": True,
        "shape": (2, 2, 1),
    },
    "A.var_fixed": {
        "value": np.ones((2, 2, 1)),
        "trainable": False,
        "shape": (2, 2, 1),
    },
}

example_module_list_variable_dict = {
    "submodule_list[0].var_trainable": example_tf_module_variable_dict["A.var_trainable"],
    "submodule_list[0].var_fixed": example_tf_module_variable_dict["A.var_fixed"],
    "submodule_list[1].var_trainable": example_tf_module_variable_dict["A.var_trainable"],
    "submodule_list[1].var_fixed": example_tf_module_variable_dict["A.var_fixed"],
    "submodule_dict['a'].var_trainable": example_tf_module_variable_dict["A.var_trainable"],
    "submodule_dict['a'].var_fixed": example_tf_module_variable_dict["A.var_fixed"],
    "submodule_dict['b'].var_trainable": example_tf_module_variable_dict["A.var_trainable"],
    "submodule_dict['b'].var_fixed": example_tf_module_variable_dict["A.var_fixed"],
    "B.var_trainable": example_tf_module_variable_dict["A.var_trainable"],
    "B.var_fixed": example_tf_module_variable_dict["A.var_fixed"],
}

kernel_param_dict = {
    "SquaredExponential.lengthscales": {
        "value": Data.ls,
        "trainable": False,
        "shape": (),
    },
    "SquaredExponential.variance": {"value": Data.var, "trainable": True, "shape": ()},
}

compose_kernel_param_dict = {
    "kernels[0].kernels[0].variance": kernel_param_dict["SquaredExponential.variance"],
    "kernels[0].kernels[0].lengthscales": kernel_param_dict["SquaredExponential.lengthscales"],
    "kernels[0].kernels[1].variance": kernel_param_dict["SquaredExponential.variance"],
    "kernels[0].kernels[1].lengthscales": kernel_param_dict["SquaredExponential.lengthscales"],
    "kernels[1].kernels[0].variance": kernel_param_dict["SquaredExponential.variance"],
    "kernels[1].kernels[0].lengthscales": kernel_param_dict["SquaredExponential.lengthscales"],
    "kernels[1].kernels[1].variance": kernel_param_dict["SquaredExponential.variance"],
    "kernels[1].kernels[1].lengthscales": kernel_param_dict["SquaredExponential.lengthscales"],
}

model_gp_param_dict = {
    "kernel.lengthscales": kernel_param_dict["SquaredExponential.lengthscales"],
    "kernel.variance": kernel_param_dict["SquaredExponential.variance"],
    "likelihood.variance": {"value": 1.0, "trainable": True, "shape": ()},
    "inducing_variable.Z": {
        "value": Data.Z,
        "trainable": True,
        "shape": (Data.M, Data.D),
    },
    "SVGP.q_mu": {
        "value": np.zeros((Data.M, 1)),
        "trainable": False,
        "shape": (Data.M, 1),
    },
    "SVGP.q_sqrt": {
        "value": np.ones((Data.M, 1)),
        "trainable": True,
        "shape": (Data.M, 1),
    },
}

example_dag_module_param_dict = {
    "SVGP.kernel.variance\nSVGP.kernel.lengthscales": kernel_param_dict[
        "SquaredExponential.lengthscales"
    ],
    "SVGP.likelihood.variance": {"value": 1.0, "trainable": True, "shape": ()},
    "SVGP.inducing_variable.Z": {
        "value": Data.Z,
        "trainable": True,
        "shape": (Data.M, Data.D),
    },
    "SVGP.q_mu": {
        "value": np.zeros((Data.M, 1)),
        "trainable": False,
        "shape": (Data.M, 1),
    },
    "SVGP.q_sqrt": {
        "value": np.ones((Data.M, 1)),
        "trainable": True,
        "shape": (Data.M, 1),
    },
}

compose_kernel_param_print_string = """\
name                                        class      transform    prior    trainable    shape    dtype      value\n\
------------------------------------------  ---------  -----------  -------  -----------  -------  -------  -------\n\
Product.kernels[0].kernels[0].variance      Parameter  Softplus              True         ()       float64        1\n\
Product.kernels[0].kernels[0].lengthscales  Parameter  Softplus              False        ()       float64        2\n\
Product.kernels[0].kernels[1].variance      Parameter  Softplus              True         ()       float64        1\n\
Product.kernels[0].kernels[1].lengthscales  Parameter  Softplus              False        ()       float64        2\n\
Product.kernels[1].kernels[0].variance      Parameter  Softplus              True         ()       float64        1\n\
Product.kernels[1].kernels[0].lengthscales  Parameter  Softplus              False        ()       float64        2\n\
Product.kernels[1].kernels[1].variance      Parameter  Softplus              True         ()       float64        1\n\
Product.kernels[1].kernels[1].lengthscales  Parameter  Softplus              False        ()       float64        2"""

kernel_param_print_string = """\
name                             class      transform    prior    trainable    shape    dtype      value\n\
-------------------------------  ---------  -----------  -------  -----------  -------  -------  -------\n\
SquaredExponential.variance      Parameter  Softplus              True         ()       float64        1\n\
SquaredExponential.lengthscales  Parameter  Softplus              False        ()       float64        2"""

kernel_param_print_string_with_shift = """\
name                             class      transform         prior    trainable    shape    dtype      value\n\
-------------------------------  ---------  ----------------  -------  -----------  -------  -------  -------\n\
SquaredExponential.variance      Parameter  Softplus + Shift           True         ()       float64        1\n\
SquaredExponential.lengthscales  Parameter  Softplus + Shift           False        ()       float64        2"""

model_gp_param_print_string = """\
name                      class      transform    prior    trainable    shape    dtype    value\n\
------------------------  ---------  -----------  -------  -----------  -------  -------  --------\n\
SVGP.kernel.variance      Parameter  Softplus              True         ()       float64  1.0\n\
SVGP.kernel.lengthscales  Parameter  Softplus              False        ()       float64  2.0\n\
SVGP.likelihood.variance  Parameter  Softplus              True         ()       float64  1.0\n\
SVGP.inducing_variable.Z  Parameter  Identity              True         (10, 1)  float64  [[0.5...\n\
SVGP.q_mu                 Parameter  Identity              False        (10, 1)  float64  [[0....\n\
SVGP.q_sqrt               Parameter  Softplus              True         (10, 1)  float64  [[1...."""

example_tf_module_variable_print_string = """\
name             class             transform    prior    trainable    shape      dtype    value\n\
---------------  ----------------  -----------  -------  -----------  ---------  -------  --------\n\
A.var_trainable  ResourceVariable                        True         (2, 2, 1)  float32  [[[0....\n\
A.var_fixed      ResourceVariable                        False        (2, 2, 1)  float32  [[[1...."""

example_module_list_variable_print_string = """\
name                                 class             transform    prior    trainable    shape      dtype    value\n\
-----------------------------------  ----------------  -----------  -------  -----------  ---------  -------  --------\n\
B.submodule_list[0].var_trainable    ResourceVariable                        True         (2, 2, 1)  float32  [[[0....\n\
B.submodule_list[0].var_fixed        ResourceVariable                        False        (2, 2, 1)  float32  [[[1....\n\
B.submodule_list[1].var_trainable    ResourceVariable                        True         (2, 2, 1)  float32  [[[0....\n\
B.submodule_list[1].var_fixed        ResourceVariable                        False        (2, 2, 1)  float32  [[[1....\n\
B.submodule_dict['a'].var_trainable  ResourceVariable                        True         (2, 2, 1)  float32  [[[0....\n\
B.submodule_dict['a'].var_fixed      ResourceVariable                        False        (2, 2, 1)  float32  [[[1....\n\
B.submodule_dict['b'].var_trainable  ResourceVariable                        True         (2, 2, 1)  float32  [[[0....\n\
B.submodule_dict['b'].var_fixed      ResourceVariable                        False        (2, 2, 1)  float32  [[[1....\n\
B.var_trainable                      ResourceVariable                        True         (2, 2, 1)  float32  [[[0....\n\
B.var_fixed                          ResourceVariable                        False        (2, 2, 1)  float32  [[[1...."""

# Note: we use grid format here because we have a double reference to the same variable
# which does not render nicely in the table formatting.
# The internal structure of the Keras model changed in TensorFlow versions 2.5.0 and 2.6.0.
example_tf_keras_model_tf_2_6_0 = """\
+-------------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+\n\
| name                          | class            | transform   | prior   | trainable   | shape     | dtype   | value    |\n\
+===============================+==================+=============+=========+=============+===========+=========+==========+\n\
| C._trainable_weights[0]       | ResourceVariable |             |         | True        | (2, 2, 1) | float32 | [[[0.... |\n\
| C.variable                    |                  |             |         |             |           |         |          |\n\
+-------------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+\n\
| C._trainable_weights[1]       | ResourceVariable |             |         | True        | ()        | float64 | 0.0      |\n\
+-------------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+\n\
| C._self_tracked_trackables[0] | Parameter        | Identity    |         | True        | ()        | float64 | 0.0      |\n\
| C.param                       |                  |             |         |             |           |         |          |\n\
+-------------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+"""


example_tf_keras_model_tf_2_5_0 = """\
+-------------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+\n\
| name                          | class            | transform   | prior   | trainable   | shape     | dtype   | value    |\n\
+===============================+==================+=============+=========+=============+===========+=========+==========+\n\
| C._trainable_weights[0]       | ResourceVariable |             |         | True        | (2, 2, 1) | float32 | [[[0.... |\n\
| C.variable                    |                  |             |         |             |           |         |          |\n\
+-------------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+\n\
| C._self_tracked_trackables[0] | Parameter        | Identity    |         | True        | ()        | float64 | 0.0      |\n\
| C.param                       |                  |             |         |             |           |         |          |\n\
+-------------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+"""

example_tf_keras_model = """\
+-------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+\n\
| name                    | class            | transform   | prior   | trainable   | shape     | dtype   | value    |\n\
+=========================+==================+=============+=========+=============+===========+=========+==========+\n\
| C._trainable_weights[0] | ResourceVariable |             |         | True        | (2, 2, 1) | float32 | [[[0.... |\n\
| C.variable              |                  |             |         |             |           |         |          |\n\
+-------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+\n\
| C.param                 | Parameter        | Identity    |         | True        | ()        | float64 | 0.0      |\n\
+-------------------------+------------------+-------------+---------+-------------+-----------+---------+----------+"""

if Version(tf.__version__) >= Version("2.6"):
    example_tf_keras_model = example_tf_keras_model_tf_2_6_0
elif Version(tf.__version__) >= Version("2.5"):
    example_tf_keras_model = example_tf_keras_model_tf_2_5_0

# ------------------------------------------
# Fixtures
# ------------------------------------------


@pytest.fixture(params=[A, B, create_kernel, create_model])
def module(request: SubRequest) -> Any:
    return request.param()


@pytest.fixture
def dag_module() -> gpflow.models.GPModel:
    dag = create_model()
    dag.kernel.variance = dag.kernel.lengthscales
    return dag


# ------------------------------------------
# Tests
# ------------------------------------------


def test_leaf_components_only_returns_parameters_and_variables(module: Any) -> None:
    for path, variable in leaf_components(module).items():
        assert isinstance(variable, tf.Variable) or isinstance(variable, gpflow.Parameter)


@pytest.mark.parametrize(
    "module_callable, expected_param_dicts",
    [(create_kernel, kernel_param_dict), (create_model, model_gp_param_dict)],
)
def test_leaf_components_registers_variable_properties(
    module_callable: Callable[[], Any], expected_param_dicts: Mapping[str, Any]
) -> None:
    module = module_callable()
    for path, variable in leaf_components(module).items():
        param_name = path.split(".")[-2] + "." + path.split(".")[-1]
        assert isinstance(variable, gpflow.Parameter)
        np.testing.assert_almost_equal(variable.numpy(), expected_param_dicts[param_name]["value"])
        assert variable.trainable == expected_param_dicts[param_name]["trainable"]
        assert variable.shape == expected_param_dicts[param_name]["shape"]


@pytest.mark.parametrize(
    "module_callable, expected_param_dicts",
    [
        (create_compose_kernel, compose_kernel_param_dict),
    ],
)
def test_leaf_components_registers_compose_kernel_variable_properties(
    module_callable: Callable[[], Any], expected_param_dicts: Mapping[str, Any]
) -> None:
    module = module_callable()
    leaf_components_dict = leaf_components(module)
    assert len(leaf_components_dict) > 0
    for path, variable in leaf_components_dict.items():
        path_as_list = path.split(".")
        param_name = path_as_list[-3] + "." + path_as_list[-2] + "." + path_as_list[-1]
        assert isinstance(variable, gpflow.Parameter)
        np.testing.assert_almost_equal(variable.numpy(), expected_param_dicts[param_name]["value"])
        assert variable.trainable == expected_param_dicts[param_name]["trainable"]
        assert variable.shape == expected_param_dicts[param_name]["shape"]


@pytest.mark.parametrize(
    "module_class, expected_var_dicts",
    [
        (A, example_tf_module_variable_dict),
        (B, example_module_list_variable_dict),
    ],
)
def test_leaf_components_registers_param_properties(
    module_class: Type[Any], expected_var_dicts: Mapping[str, Any]
) -> None:
    module = module_class()
    for path, variable in leaf_components(module).items():
        var_name = path.split(".")[-2] + "." + path.split(".")[-1]
        assert isinstance(variable, tf.Variable)
        np.testing.assert_equal(variable.numpy(), expected_var_dicts[var_name]["value"])
        assert variable.trainable == expected_var_dicts[var_name]["trainable"]
        assert variable.shape == expected_var_dicts[var_name]["shape"]


@pytest.mark.parametrize("expected_var_dicts", [example_dag_module_param_dict])
def test_merge_leaf_components_merges_keys_with_same_values(
    dag_module: Any, expected_var_dicts: Mapping[str, Any]
) -> None:
    leaf_components_dict = leaf_components(dag_module)
    for path, variable in _merge_leaf_components(leaf_components_dict).items():
        assert path in expected_var_dicts
        for sub_path in path.split("\n"):
            assert sub_path in leaf_components_dict
            assert leaf_components_dict[sub_path] is variable


@pytest.mark.parametrize(
    "module_callable, expected_param_print_string",
    [
        (create_compose_kernel, compose_kernel_param_print_string),
        (create_kernel, kernel_param_print_string),
        (create_model, model_gp_param_print_string),
        (A, example_tf_module_variable_print_string),
        (B, example_module_list_variable_print_string),
    ],
)
def test_print_summary_output_string(
    module_callable: Callable[[], Any], expected_param_print_string: str
) -> None:
    with as_context(Config(positive_minimum=0.0)):
        assert tabulate_module_summary(module_callable()) == expected_param_print_string


def test_print_summary_output_string_with_positive_minimum() -> None:
    with as_context(Config(positive_minimum=1e-6)):
        assert tabulate_module_summary(create_kernel()) == kernel_param_print_string_with_shift


def test_print_summary_for_keras_model() -> None:
    # Note: best to use `grid` formatting for `tf.keras.Model` printing
    # because of the duplicates in the references to the variables.
    assert tabulate_module_summary(C(), tablefmt="grid") == example_tf_keras_model


def test_leaf_components_combination_kernel() -> None:
    """
    Regression test for kernel compositions - output for printing should not be empty (issue #1066).
    """
    k = gpflow.kernels.SquaredExponential() + gpflow.kernels.SquaredExponential()
    assert leaf_components(k), "Combination kernel should have non-empty leaf components"


def test_module_parameters_return_iterators_not_generators() -> None:
    """
    Regression test: Ensure that gpflow.Module parameters return iterators like in TF2, not
    generators.

    Reason:
    param = m.params  # <generator object>
    x = [p for p in param] # List[Parameters]
    y = [p for p in param] # [] empty!
    """
    m = create_model()
    assert isinstance(m, gpflow.base.Module)
    assert isinstance(m.parameters, tuple)
    assert isinstance(m.trainable_parameters, tuple)


# tests/gpflow/utilities/test_utilities.py

from typing import Tuple

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

import gpflow
from gpflow.base import AnyNDArray
from gpflow.models.util import data_input_to_tensor


def test_select_parameters_with_prior() -> None:
    kernel = gpflow.kernels.SquaredExponential()
    params = gpflow.utilities.select_dict_parameters_with_prior(kernel)
    assert params == {}

    kernel.variance.prior = tfp.distributions.Gamma(1.0, 1.0)
    params = gpflow.utilities.select_dict_parameters_with_prior(kernel)
    assert len(params) == 1


def test_data_input_to_tensor() -> None:
    input1 = (1.0, (2.0,))
    output1 = data_input_to_tensor(input1)
    assert output1[0].dtype == tf.float64
    assert output1[1][0].dtype == tf.float64

    input2 = (1.0, [2.0])
    output2 = data_input_to_tensor(input2)
    assert output2[0].dtype == tf.float64
    assert output2[1][0].dtype == tf.float64

    input3: Tuple[float, Tuple[AnyNDArray, AnyNDArray]] = (
        1.0,
        (np.arange(3, dtype=np.float16),) * 2,
    )
    output3 = data_input_to_tensor(input3)
    assert output3[0].dtype == tf.float64
    assert output3[1][0].dtype == tf.float16
    assert output3[1][1].dtype == tf.float16


# doc/build_docs.py

# Copyright 2022 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.
"""
Code for building GPflow's documentation for a specified branch.
"""
import argparse
import shutil
import subprocess
from itertools import chain
from pathlib import Path
from time import perf_counter
from typing import Collection, Optional

from generate_module_rst import generate_module_rst
from tabulate import tabulate
from versions import Branch

import gpflow

_SRC = Path(__file__).parent
_SPHINX_SRC = _SRC / "sphinx"
_NOTEBOOKS_SRC = _SPHINX_SRC / "notebooks"

_TMP = Path("/tmp/gpflow_build_docs")
_BUILD_TMP = _TMP / "build"
_NOTEBOOKS_TMP = _BUILD_TMP / "notebooks"
_DOCTREE_TMP = _TMP / "doctree"


def _create_fake_notebook(
    destination_relative_path: Path, limit_notebooks: Collection[str]
) -> None:
    limiting_command = f"--limit_notebooks {' '.join(limit_notebooks)}"
    print(f"Generating fake, due to {limiting_command}")

    destination = _NOTEBOOKS_TMP / destination_relative_path
    title = f"Fake {destination.name}"
    title_line = "#" * len(title)

    destination.write_text(
        f"""{title}
{title_line}

Fake {destination.name} due to::

   {limiting_command}
"""
    )


def _build_notebooks(limit_notebooks: Optional[Collection[str]]) -> None:
    # Building the notebooks is really slow. Let's time it so we know which notebooks we can /
    # should optimise.
    timings = []
    all_notebooks = chain(_NOTEBOOKS_TMP.glob("**/*.pct.py"), _NOTEBOOKS_TMP.glob("**/*.md"))
    for i, source_path in enumerate(all_notebooks):
        before = perf_counter()
        print()
        print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")
        print("Building:", source_path)
        print("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX")

        source_relative_path = source_path.relative_to(_NOTEBOOKS_TMP)
        destination_relative_path = source_relative_path
        while destination_relative_path.suffix:  # .pct.py has several suffixes. Remove all of them.
            destination_relative_path = destination_relative_path.with_suffix("")
        destination_relative_path = destination_relative_path.with_suffix(".ipynb")

        if limit_notebooks is None or destination_relative_path.stem in limit_notebooks:
            subprocess.run(
                [
                    "jupytext",
                    "--execute",
                    "--to",
                    "notebook",
                    "-o",
                    str(destination_relative_path),
                    str(source_relative_path),
                ],
                cwd=_NOTEBOOKS_TMP,
            ).check_returncode()
        else:
            _create_fake_notebook(destination_relative_path, limit_notebooks)

        after = perf_counter()
        timings.append((after - before, source_relative_path))

    timings.sort(reverse=True)
    print()
    print("Notebooks by build-time:")
    print(tabulate(timings, headers=["Time", "Notebook"]))
    print()


def main() -> None:
    parser = argparse.ArgumentParser(description="Build the GPflow documentation.")
    parser.add_argument(
        "branch",
        type=str,
        choices=[b.value for b in Branch],
        help="Git branch that is currently being built.",
    )
    parser.add_argument(
        "destination",
        type=Path,
        help="Directory to write docs to.",
    )
    parser.add_argument(
        "--limit_notebooks",
        "--limit-notebooks",
        type=str,
        nargs="*",
        help="Only process the notebooks with this base/stem name. Useful when debugging.",
    )
    parser.add_argument(
        "--fail_on_warning",
        "--fail-on-warning",
        default=False,
        action="store_true",
        help="If set, crash if there were any warnings while generating documentation.",
    )
    args = parser.parse_args()
    branch = Branch(args.branch)
    dest = args.destination
    version_dest = dest / branch.version

    shutil.rmtree(version_dest, ignore_errors=True)
    shutil.rmtree(_TMP, ignore_errors=True)
    _TMP.mkdir(parents=True)

    shutil.copytree(_SPHINX_SRC, _BUILD_TMP)
    (_BUILD_TMP / "build_version.txt").write_text(branch.version)
    _build_notebooks(args.limit_notebooks)
    generate_module_rst(gpflow, _BUILD_TMP / "api")

    sphinx_commands = [
        "sphinx-build",
        "-b",
        "html",
        "-d",
        str(_DOCTREE_TMP),
        str(_BUILD_TMP),
        str(version_dest),
    ]
    if args.fail_on_warning:
        sphinx_commands.extend(
            [
                "-W",
                "--keep-going",
            ]
        )

    subprocess.run(sphinx_commands).check_returncode()


if __name__ == "__main__":
    main()


# doc/generate_module_rst.py

# Copyright 2019 GPflow Authors
#
# 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.
"""Script to autogenerate .rst files for autodocumentation of classes and modules in GPflow.
To be run by the CI system to update docs.
"""
import inspect
from dataclasses import dataclass
from pathlib import Path
from types import ModuleType
from typing import Any, Callable, Deque, Dict, List, Mapping, Optional, Set, TextIO, Type, Union

from gpflow.utilities import Dispatcher

RST_LEVEL_SYMBOLS = ["=", "-", "~", '"', "'", "^"]

IGNORE_MODULES = {
    "gpflow.covariances.dispatch",
    "gpflow.conditionals.dispatch",
    "gpflow.expectations.dispatch",
    "gpflow.kullback_leiblers.dispatch",
    "gpflow.versions",
}


def _header(header: str, level: int) -> str:
    return f"{header}\n{RST_LEVEL_SYMBOLS[level] * len(header)}"


@dataclass
class DocumentableDispatcher:

    name: str
    obj: Dispatcher

    def implementations(self) -> Mapping[Callable[..., Any], List[Type[Any]]]:
        implementations: Dict[Callable[..., Any], List[Type[Any]]] = {}
        for args, impl in self.obj.funcs.items():
            implementations.setdefault(impl, []).append(args)
        return implementations

    def write(self, out: TextIO) -> None:
        out.write(
            f"""
{_header(self.name, 2)}

This function uses multiple dispatch, which will depend on the type of argument passed in:
"""
        )
        for impl, argss in self.implementations().items():
            impl_name = f"{impl.__module__}.{impl.__name__}"

            out.write(
                """
.. code-block:: python

"""
            )
            for args in argss:
                arg_names = ", ".join([a.__name__ for a in args])
                out.write(f"    {self.name}( {arg_names} )\n")
            out.write(f"    # dispatch to -> {impl_name}(...)\n")
            out.write(
                f"""
.. autofunction:: {impl_name}
"""
            )


@dataclass
class DocumentableClass:

    name: str
    obj: Type[Any]

    def write(self, out: TextIO) -> None:
        out.write(
            f"""
{_header(self.name, 2)}

.. autoclass:: {self.name}
   :show-inheritance:
   :members:
"""
        )


@dataclass
class DocumentableFunction:

    name: str
    obj: Callable[..., Any]

    def write(self, out: TextIO) -> None:
        out.write(
            f"""
{_header(self.name, 2)}

.. autofunction:: {self.name}
"""
        )


@dataclass
class DocumentableModule:

    name: str
    obj: ModuleType
    modules: List["DocumentableModule"]
    classes: List[DocumentableClass]
    functions: List[Union[DocumentableDispatcher, DocumentableFunction]]

    @staticmethod
    def collect(
        root: ModuleType,
    ) -> "DocumentableModule":
        root_name = root.__name__
        exported_names = set(getattr(root, "__all__", []))

        modules: List["DocumentableModule"] = []
        classes: List[DocumentableClass] = []
        functions: List[Union[DocumentableDispatcher, DocumentableFunction]] = []

        for key in dir(root):
            if key.startswith("_"):
                continue

            child = getattr(root, key)
            child_name = root_name + "." + key
            if child_name in IGNORE_MODULES:
                continue

            def _should_ignore(child: Union[Callable[..., Any], Type[Any]]) -> bool:
                declared_in_root = child.__module__ == root_name
                explicitly_exported = key in exported_names
                return not (declared_in_root or explicitly_exported)

            if isinstance(child, Dispatcher):
                functions.append(DocumentableDispatcher(child_name, child))
            elif inspect.ismodule(child):
                if child.__name__ != child_name:  # Ignore imports of modules.
                    continue
                modules.append(DocumentableModule.collect(child))
            elif inspect.isclass(child):
                if _should_ignore(child):
                    continue
                classes.append(DocumentableClass(child_name, child))
            elif inspect.isfunction(child):
                if _should_ignore(child):
                    continue
                functions.append(DocumentableFunction(child_name, child))

        return DocumentableModule(root_name, root, modules, classes, functions)

    def seen_in_dispatchers(self, seen: Set[int]) -> None:
        for module in self.modules:
            module.seen_in_dispatchers(seen)
        for function in self.functions:
            if isinstance(function, DocumentableDispatcher):
                impls = function.obj.funcs.values()
                for impl in impls:
                    seen.add(id(impl))

    def prune_duplicates(self) -> None:
        seen: Set[int] = set()
        self.seen_in_dispatchers(seen)

        # Breadth-first search so that we prefer objects with shorter names.
        todo = Deque([self])
        while todo:
            module = todo.popleft()

            new_classes = []
            for c in module.classes:
                if id(c.obj) not in seen:
                    seen.add(id(c.obj))
                    new_classes.append(c)
            module.classes = new_classes

            new_functions = []
            for f in module.functions:
                if id(f.obj) not in seen:
                    seen.add(id(f.obj))
                    new_functions.append(f)
            module.functions = new_functions

            todo.extend(module.modules)

    def prune_empty_modules(self) -> None:
        new_modules = []
        for m in self.modules:
            m.prune_empty_modules()

            if m.modules or m.classes or m.functions:
                new_modules.append(m)
        self.modules = new_modules

    def prune(self) -> None:
        self.prune_duplicates()
        self.prune_empty_modules()

    def write_modules(self, out: TextIO) -> None:
        if not self.modules:
            return

        out.write(
            f"""
{_header('Modules', 1)}

.. toctree::
   :maxdepth: 1

"""
        )
        for module in self.modules:
            out.write(f"   {module.name} <{module.name.split('.')[-1]}/index>\n")

    def write_classes(self, out: TextIO) -> None:
        if not self.classes:
            return

        out.write(
            f"""
{_header('Classes', 1)}
"""
        )
        for cls in self.classes:
            cls.write(out)

    def write_functions(self, out: TextIO) -> None:
        if not self.functions:
            return

        out.write(
            f"""
{_header('Functions', 1)}
"""
        )
        for function in self.functions:
            function.write(out)

    def write(self, path: Path) -> None:
        dir_path = path / f"{self.name.replace('.', '/')}"
        dir_path.mkdir(parents=True, exist_ok=True)
        index_path = dir_path / "index.rst"
        with index_path.open("wt") as out:
            print("Writing", index_path)
            out.write(
                f"""{_header(self.name, 0)}

.. THIS IS AN AUTOGENERATED RST FILE

.. automodule:: {self.name}
"""
            )
            self.write_modules(out)
            self.write_classes(out)
            self.write_functions(out)

        for module in self.modules:
            module.write(path)

    def str_into(self, indent: int, lines: List[str]) -> None:
        lines.append(2 * indent * " " + "Module: " + self.name)
        for module in self.modules:
            module.str_into(indent + 1, lines)
        for cls in self.classes:
            lines.append(2 * (indent + 1) * " " + "Class: " + cls.name)
        for function in self.functions:
            lines.append(2 * (indent + 1) * " " + "Function: " + function.name)

    def __str__(self) -> str:
        lines: List[str] = []
        self.str_into(0, lines)
        return "\n".join(lines)


def generate_module_rst(module: ModuleType, dest: Path) -> None:
    """
    Traverses the given `module` and generates `.rst` files for Sphinx.
    """
    docs = DocumentableModule.collect(module)
    docs.prune()
    docs.write(dest)


# doc/update_versions.py

# Copyright 2022 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.
"""
Tool for updating the `versions.json` and `index.html` in our documentation.
"""
import argparse
import json
from pathlib import Path
from typing import Any, List, Mapping, Optional

from packaging.version import Version
from versions import Branch


class _Versions:
    def __init__(self, versions: List[Mapping[str, Any]]) -> None:
        self._versions = versions
        self._sorted = False

    @staticmethod
    def _versions_json(dest: Path) -> Path:
        return dest / "versions.json"

    @staticmethod
    def read(docs_dir: Path) -> "_Versions":
        if _Versions._versions_json(docs_dir).exists():
            with _Versions._versions_json(docs_dir).open("rt") as versions_file:
                versions = json.load(versions_file)
        else:
            versions = {}
        return _Versions(versions)

    def write(self, docs_dir: Path) -> None:
        with _Versions._versions_json(docs_dir).open("wt") as versions_file:
            json.dump(self._versions, versions_file, indent=2)

    def add(self, version: str) -> None:
        self._versions = [v for v in self._versions if v["version"] != version]
        self._versions.append(
            {
                "version": version,
                "url": f"https://gpflow.github.io/GPflow/{version}/index.html",
            }
        )
        self._sorted = False

    def sort(self) -> None:
        develop: Optional[Mapping[str, Any]] = None
        for v in self._versions:
            if v["version"] == "develop":
                develop = v
        versions = [v for v in self._versions if v["version"] != "develop"]
        versions.sort(key=lambda v: Version(v["version"]), reverse=True)
        if develop is not None:
            # Insert `develop` in the second spot, after the latest release, but before older
            # releases:
            versions.insert(1, develop)
        self._versions = versions
        self._sorted = True

    @property
    def latest(self) -> str:
        assert self._sorted, "Must sort versions to find latest."
        latest = self._versions[0]["version"]
        assert isinstance(latest, str)  # Hint for mypy.
        return latest


def _create_root_redirect(versions: _Versions, dest: Path) -> None:

    (dest / "index.html").write_text(
        f"""<html xmlns="http://www.w3.org/1999/xhtml">
  <head>
    <title>GPflow documentation</title>
    <meta http-equiv="refresh" content="0;URL='{versions.latest}/index.html'" />
  </head>
  <body>
    <p>Redirecting to the most recent release.</p>
  </body>
</html>
"""
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Build the GPflow documentation.")
    parser.add_argument(
        "branch",
        type=str,
        choices=[b.value for b in Branch],
        help="Git branch that is currently being built.",
    )
    parser.add_argument("docs_dir", type=Path, help="To read / write docs versions from.")
    args = parser.parse_args()
    branch = Branch(args.branch)
    docs_dir = args.docs_dir

    versions = _Versions.read(docs_dir)
    versions.add(branch.version)
    versions.sort()
    versions.write(docs_dir)

    _create_root_redirect(versions, docs_dir)


if __name__ == "__main__":
    main()


# doc/versions.py

# Copyright 2022 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.
"""
Code for keeping track of different versions of documentation.
"""
from enum import Enum
from pathlib import Path

# gpflow.__version__ does not seem reliable when used from a branch, so we read the VERSION file
# directly. (That's also a lot faster than importing gpflow...)
_GPFLOW_VERSION_FILE = Path(__file__).parent.parent / "VERSION"


class Branch(Enum):
    DEVELOP = "develop"
    """
    Building the `develop` branch means we're creating temporary documentation.

    Use this if in doubt.
    """

    MASTER = "master"
    """
    Building the `master` branch means we're creating a new release.
    """

    @property
    def version(self) -> str:
        """
        Returns the version of code / documentation this branch refers to.
        """
        if self == Branch.DEVELOP:
            return "develop"
        else:
            assert self == Branch.MASTER
            return _GPFLOW_VERSION_FILE.read_text().strip()


# gpflow/__init__.py

# Copyright 2016-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.

# flake8: noqa

from . import (
    conditionals,
    config,
    covariances,
    expectations,
    experimental,
    functions,
    inducing_variables,
    kernels,
    kullback_leiblers,
    likelihoods,
    logdensities,
    mean_functions,
    models,
    monitor,
    optimizers,
    posteriors,
    probability_distributions,
    quadrature,
    utilities,
)
from .base import Module, Parameter
from .config import default_float, default_int, default_jitter
from .utilities import set_trainable
from .versions import __version__

__all__ = [
    "Module",
    "Parameter",
    "__version__",
    "base",
    "ci_utils",
    "conditionals",
    "config",
    "covariances",
    "default_float",
    "default_int",
    "default_jitter",
    "expectations",
    "experimental",
    "functions",
    "inducing_variables",
    "kernels",
    "kullback_leiblers",
    "likelihoods",
    "logdensities",
    "mean_functions",
    "models",
    "monitor",
    "mypy_flags",
    "optimizers",
    "posteriors",
    "probability_distributions",
    "quadrature",
    "set_trainable",
    "type_flags",
    "utilities",
    "versions",
]


# gpflow/base.py

# 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


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)


class PriorOn(Enum):
    CONSTRAINED = "constrained"
    UNCONSTRAINED = "unconstrained"


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

    @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]

    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


# gpflow/ci_utils.py

# 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.

# pylint: skip-file

import os
from typing import Any, Iterable, List, Sequence, Type, TypeVar


def is_continuous_integration() -> bool:
    """
    Determines whether we are running on the Continuous Integration system for
    notebook integration tests. This is used to speed up notebook integration
    tests (built on every pull request commit) by capping all expensive loops
    at a small number, rather than running until convergence. When building the
    docs (indicated by the presence of the `DOCS` environment variable), we
    need to run notebooks to completion, and this function returns `False`.
    Whether we are running on CI is determined by the presence of the `CI`
    environment variable.
    """
    if "DOCS" in os.environ:
        return False

    return "CI" in os.environ


def reduce_in_tests(n: int, test_n: int = 2) -> int:
    return test_n if is_continuous_integration() else n


def subclasses(cls: Type[Any]) -> Iterable[Type[Any]]:
    """
    Generator that returns all (not just direct) subclasses of `cls`
    """
    for subclass in cls.__subclasses__():
        yield from subclasses(subclass)
        yield subclass


# gpflow/functions.py

# Copyright 2022 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.
"""
Throughout GPflow, by default, latent functions being modelled with Gaussian
processes are assumed to have zero mean, f ~ GP(0, k(x,x')).

In some cases we may wish to model only the deviation from a fixed function
with a Gaussian process.  For flexibility this fixed function could be both
input dependent and parameterised function, μ(x; θ),
with some unknown parameters θ, resulting in f ~ GP(μ(x;θ), k(x,x')).

The GPflow :class:`MeanFunction <gpflow.mean_functions.MeanFunction>` class
allows this to be done whilst additionally learning parameters of the
parametric function.
"""
from typing import Collection, Iterator, Optional, Tuple

import numpy as np
import tensorflow as tf

from .base import Module, Parameter, TensorType
from .config import default_float, default_int
from .experimental.check_shapes import check_shape as cs
from .experimental.check_shapes import check_shapes, inherit_check_shapes


class Function(Module):
    """
    The base function class.
    To implement a function, write the __call__ method. This takes a
    tensor X and returns a tensor f(X). In accordance with the GPflow
    standard, each row of X represents one datum, and each row of Y is computed
    independently for each row of X.

    :class:`Function` classes can have parameters, see the :class:`Linear` class for an
    example.
    """

    @check_shapes(
        "X: [batch..., D]",
        "return: [batch..., Q]",
    )
    def __call__(self, X: TensorType) -> tf.Tensor:
        raise NotImplementedError("Implement the __call__ method for this mean function")

    def __add__(self, other: "Function") -> "Function":
        return Additive(self, other)

    def __mul__(self, other: "Function") -> "Function":
        return Product(self, other)


class MeanFunction(Function):
    """
    Mixin for :class:`Function`\ s that are appropriate for use as mean functions.
    """


class Additive(MeanFunction, Function):
    def __init__(self, first_part: Function, second_part: Function) -> None:
        MeanFunction.__init__(self)
        self.add_1 = first_part
        self.add_2 = second_part

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return tf.add(self.add_1(X), self.add_2(X))


class Product(MeanFunction, Function):
    def __init__(self, first_part: Function, second_part: Function):
        MeanFunction.__init__(self)

        self.prod_1 = first_part
        self.prod_2 = second_part

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return tf.multiply(self.prod_1(X), self.prod_2(X))


class Linear(MeanFunction, Function):
    """
    y_i = A x_i + b
    """

    @check_shapes(
        "A: [broadcast D, broadcast Q]",
        "b: [broadcast Q]",
    )
    def __init__(self, A: TensorType = None, b: TensorType = None) -> None:
        """
        A is a matrix which maps each element of X to Y, b is an additive
        constant.
        """
        MeanFunction.__init__(self)
        A = np.ones((1, 1), dtype=default_float()) if A is None else A
        b = np.zeros(1, dtype=default_float()) if b is None else b
        self.A = Parameter(np.atleast_2d(A))
        self.b = Parameter(b)

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return tf.tensordot(X, self.A, [[-1], [0]]) + self.b


class Identity(Linear, Function):
    """
    y_i = x_i
    """

    # The many type-ignores in this class is because we replace a field in the super class with a
    # property, which mypy doesn't like.

    def __init__(self, input_dim: Optional[int] = None) -> None:
        Linear.__init__(self)
        self.input_dim = input_dim

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        return X

    @property
    def A(self) -> tf.Tensor:  # type: ignore[override]
        if self.input_dim is None:
            raise ValueError(
                "An input_dim needs to be specified when using the "
                "`Identity` mean function in combination with expectations."
            )
        return tf.eye(self.input_dim, dtype=default_float())

    @property
    def b(self) -> tf.Tensor:  # type: ignore[override]
        if self.input_dim is None:
            raise ValueError(
                "An input_dim needs to be specified when using the "
                "`Identity` mean function in combination with expectations."
            )

        return tf.zeros(self.input_dim, dtype=default_float())

    @A.setter  # type: ignore[attr-defined, no-redef]
    def A(self, A: tf.Tensor) -> None:
        pass

    @b.setter  # type: ignore[attr-defined, no-redef]
    def b(self, b: tf.Tensor) -> None:
        pass


class Constant(MeanFunction, Function):
    @check_shapes(
        "c: [broadcast Q]",
    )
    def __init__(self, c: TensorType = None) -> None:
        super().__init__()
        c = np.zeros(1) if c is None else c
        self.c = Parameter(c)

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        tile_shape = tf.concat(
            [tf.shape(X)[:-1], [1]],
            axis=0,
        )
        reshape_shape = tf.concat(
            [tf.ones(shape=(tf.rank(X) - 1), dtype=default_int()), [-1]],
            axis=0,
        )
        return tf.tile(tf.reshape(self.c, reshape_shape), tile_shape)


class Zero(Constant, Function):
    def __init__(self, output_dim: int = 1) -> None:
        Constant.__init__(self)
        self.output_dim = output_dim
        del self.c

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        output_shape = tf.concat([tf.shape(X)[:-1], [self.output_dim]], axis=0)
        return tf.zeros(output_shape, dtype=X.dtype)


class Polynomial(MeanFunction, Function):
    """
    A generic polynomial mean function.
    """

    @check_shapes("w: [broadcast output_dim, broadcast n_terms]")
    def __init__(
        self, degree: int, input_dim: int = 1, output_dim: int = 1, w: Optional[TensorType] = None
    ) -> None:
        """
        :param degree: The degree of the polynomial.
        :param input_dim: Number of inputs / variables this polynomial is defined over.
        :param output_dim: Number of outputs / polynomials.
        :param w: Initial weights of the terms of the polynomial. The inner dimension (``n_terms``)
            should correspond to the powers returned by ``compute_powers``.
        """
        powers = cs(tuple(self.compute_powers(degree, input_dim)), "[n_terms, input_dim]")
        if w is None:
            w = cs([1.0] + (len(powers) - 1) * [0.0], "[n_terms]")
        w_shape = (output_dim, len(powers))
        self.powers = tf.constant(powers, dtype=default_float())
        self.w = Parameter(tf.broadcast_to(w, w_shape))

    @staticmethod
    def compute_powers(degree: int, input_dim: int) -> Iterator[Tuple[int, ...]]:
        """
        Computes integer tuples corresponding to the powers to raise inputs to.

        Specifically this returns, in lexicographical order, all tuples where:
        * The tuple has length `input_dim`.
        * The values are non-negative integers.
        * The sum of the tuple is no greater than `degree`.

        For example::

            compute_powers(degree=2, input_dim=3)

        returns::

            (0, 0, 0)
            (0, 0, 1)
            (0, 0, 2)
            (0, 1, 0)
            (0, 1, 1)
            (0, 2, 0)
            (1, 0, 0)
            (1, 0, 1)
            (1, 1, 0)
            (2, 0, 0)

        where a tuple::

            (1, 0, 2)

        will translate to a the term::

            w[i] * (x[0]**1) * (x[1]**0) * (x[2]**2)
        """
        if not input_dim:
            yield ()
            return
        for i in range(degree + 1):
            for inner in Polynomial.compute_powers(degree - i, input_dim - 1):
                yield (i,) + inner

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        raised = cs(tf.pow(X[..., None, :], self.powers), "[batch..., n_terms, input_dim]")
        prod = cs(tf.math.reduce_prod(raised, axis=-1), "[batch..., n_terms]")
        return tf.einsum("...i,ji->...j", prod, self.w)


class SwitchedFunction(MeanFunction, Function):
    """
    This class enables to use different (independent) functions respective
    to the data 'label'.
    We assume the 'label' is stored in the extra column of X.
    """

    def __init__(self, function_list: Collection[Function]) -> None:
        super().__init__()
        self.functions = function_list

    @inherit_check_shapes
    def __call__(self, X: TensorType) -> tf.Tensor:
        ind = tf.gather(tf.transpose(X), tf.shape(X)[1] - 1)  # ind = X[:,-1]
        ind = tf.cast(ind, tf.int32)
        X = tf.transpose(
            tf.gather(tf.transpose(X), tf.range(0, tf.shape(X)[1] - 1))
        )  # X = X[:,:-1]

        # split up X into chunks corresponding to the relevant likelihoods
        x_list = tf.dynamic_partition(X, ind, len(self.meanfunctions))
        # apply the likelihood-function to each section of the data
        results = [m(x) for x, m in zip(x_list, self.meanfunctions)]
        # stitch the results back together
        partitions = tf.dynamic_partition(tf.range(0, tf.size(ind)), ind, len(self.meanfunctions))
        return tf.dynamic_stitch(partitions, results)


class SwitchedMeanFunction(SwitchedFunction):
    """
    Renamed :class:`SwitchedFunction` for backwards compatibility.
    """

    def __init__(self, meanfunction_list: Collection[MeanFunction]) -> None:
        super().__init__(function_list=meanfunction_list)

    @property
    def meanfunctions(self) -> Collection[MeanFunction]:
        return self.functions

    @meanfunctions.setter
    def meanfunctions(self, value: Collection[MeanFunction]) -> None:
        self.function_list = value


# gpflow/kullback_leiblers.py

# Copyright 2016-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.

# -*- coding: utf-8 -*-

import tensorflow as tf
from packaging.version import Version

from .base import TensorType
from .config import default_float, default_jitter
from .covariances import Kuu
from .experimental.check_shapes import check_shapes
from .inducing_variables import InducingVariables
from .kernels import Kernel
from .utilities import Dispatcher, to_default_float

prior_kl = Dispatcher("prior_kl")


@prior_kl.register(InducingVariables, Kernel, object, object)
@check_shapes(
    "inducing_variable: [N, D, broadcast L]",
    "q_mu: [M, L]",
    "q_sqrt: [M_L_or_L_M_M...]",
    "return: []",
)
def _(
    inducing_variable: InducingVariables,
    kernel: Kernel,
    q_mu: TensorType,
    q_sqrt: TensorType,
    whiten: bool = False,
) -> tf.Tensor:
    if whiten:
        return gauss_kl(q_mu, q_sqrt, None)
    else:
        K = Kuu(inducing_variable, kernel, jitter=default_jitter())  # [P, M, M] or [M, M]
        return gauss_kl(q_mu, q_sqrt, K)


@check_shapes(
    "q_mu: [M, L]",
    "q_sqrt: [M_L_or_L_M_M...]",
    "K: [broadcast L, M, M]",
    "K_cholesky: [broadcast L, M, M]",
    "return: []",
)
def gauss_kl(
    q_mu: TensorType, q_sqrt: TensorType, K: TensorType = None, *, K_cholesky: TensorType = None
) -> tf.Tensor:
    """
    Compute the KL divergence KL[q || p] between::

          q(x) = N(q_mu, q_sqrt^2)

    and::

          p(x) = N(0, K)    if K is not None
          p(x) = N(0, I)    if K is None

    We assume L multiple independent distributions, given by the columns of
    q_mu and the first or last dimension of q_sqrt. Returns the *sum* of the
    divergences.

    q_mu is a matrix ([M, L]), each column contains a mean.

    - q_sqrt can be a 3D tensor ([L, M, M]), each matrix within is a lower
      triangular square-root matrix of the covariance of q.
    - q_sqrt can be a matrix ([M, L]), each column represents the diagonal of a
      square-root matrix of the covariance of q.

    K is the covariance of p (positive-definite matrix).  The K matrix can be
    passed either directly as `K`, or as its Cholesky factor, `K_cholesky`.  In
    either case, it can be a single matrix [M, M], in which case the sum of the
    L KL divergences is computed by broadcasting, or L different covariances
    [L, M, M].

    Note: if no K matrix is given (both `K` and `K_cholesky` are None),
    `gauss_kl` computes the KL divergence from p(x) = N(0, I) instead.
    """

    if (K is not None) and (K_cholesky is not None):
        raise ValueError(
            "Ambiguous arguments: gauss_kl() must only be passed one of `K` or `K_cholesky`."
        )

    is_white = (K is None) and (K_cholesky is None)
    is_diag = len(q_sqrt.shape) == 2

    M, L = tf.shape(q_mu)[0], tf.shape(q_mu)[1]

    if is_white:
        alpha = q_mu  # [M, L]
    else:
        if K is not None:
            Lp = tf.linalg.cholesky(K)  # [L, M, M] or [M, M]
        elif K_cholesky is not None:
            Lp = K_cholesky  # [L, M, M] or [M, M]

        is_batched = len(Lp.shape) == 3

        q_mu = tf.transpose(q_mu)[:, :, None] if is_batched else q_mu  # [L, M, 1] or [M, L]
        alpha = tf.linalg.triangular_solve(Lp, q_mu, lower=True)  # [L, M, 1] or [M, L]

    if is_diag:
        Lq = Lq_diag = q_sqrt
        Lq_full = tf.linalg.diag(tf.transpose(q_sqrt))  # [L, M, M]
    else:
        Lq = Lq_full = tf.linalg.band_part(q_sqrt, -1, 0)  # force lower triangle # [L, M, M]
        Lq_diag = tf.linalg.diag_part(Lq)  # [M, L]

    # Mahalanobis term: μqᵀ Σp⁻¹ μq
    mahalanobis = tf.reduce_sum(tf.square(alpha))

    # Constant term: - L * M
    constant = -to_default_float(tf.size(q_mu, out_type=tf.int64))

    # Log-determinant of the covariance of q(x):
    logdet_qcov = tf.reduce_sum(tf.math.log(tf.square(Lq_diag)))

    # Trace term: tr(Σp⁻¹ Σq)
    if is_white:
        trace = tf.reduce_sum(tf.square(Lq))
    else:
        if is_diag and not is_batched:
            # K is [M, M] and q_sqrt is [M, L]: fast specialisation
            LpT = tf.transpose(Lp)  # [M, M]
            Lp_inv = tf.linalg.triangular_solve(
                Lp, tf.eye(M, dtype=default_float()), lower=True
            )  # [M, M]
            K_inv = tf.linalg.diag_part(tf.linalg.triangular_solve(LpT, Lp_inv, lower=False))[
                :, None
            ]  # [M, M] -> [M, 1]
            trace = tf.reduce_sum(K_inv * tf.square(q_sqrt))
        else:
            if is_batched or Version(tf.__version__) >= Version("2.2"):
                Lp_full = Lp
            else:
                # workaround for segfaults when broadcasting in TensorFlow<2.2
                Lp_full = tf.tile(tf.expand_dims(Lp, 0), [L, 1, 1])
            LpiLq = tf.linalg.triangular_solve(Lp_full, Lq_full, lower=True)
            trace = tf.reduce_sum(tf.square(LpiLq))

    twoKL = mahalanobis + constant - logdet_qcov + trace

    # Log-determinant of the covariance of p(x):
    if not is_white:
        log_sqdiag_Lp = tf.math.log(tf.square(tf.linalg.diag_part(Lp)))
        sum_log_sqdiag_Lp = tf.reduce_sum(log_sqdiag_Lp)
        # If K is [L, M, M], num_latent_gps is no longer implicit, no need to multiply the single kernel logdet
        scale = 1.0 if is_batched else to_default_float(L)
        twoKL += scale * sum_log_sqdiag_Lp

    return 0.5 * twoKL


# gpflow/logdensities.py

# Copyright 2016-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.

import numpy as np
import tensorflow as tf

from .base import TensorType
from .experimental.check_shapes import check_shapes
from .utilities import to_default_float


@check_shapes(
    "x: [broadcast shape...]",
    "mu: [broadcast shape...]",
    "var: [broadcast shape...]",
    "return: [shape...]",
)
def gaussian(x: TensorType, mu: TensorType, var: TensorType) -> tf.Tensor:
    return -0.5 * (np.log(2 * np.pi) + tf.math.log(var) + tf.square(mu - x) / var)


@check_shapes(
    "x: [broadcast shape...]",
    "mu: [broadcast shape...]",
    "var: [broadcast shape...]",
    "return: [shape...]",
)
def lognormal(x: TensorType, mu: TensorType, var: TensorType) -> tf.Tensor:
    lnx = tf.math.log(x)
    return gaussian(lnx, mu, var) - lnx


@check_shapes(
    "x: [broadcast shape...]",
    "p: [broadcast shape...]",
    "return: [shape...]",
)
def bernoulli(x: TensorType, p: TensorType) -> tf.Tensor:
    return tf.math.log(tf.where(tf.equal(x, 1), p, 1 - p))


@check_shapes(
    "x: [broadcast shape...]",
    "lam: [broadcast shape...]",
    "return: [shape...]",
)
def poisson(x: TensorType, lam: TensorType) -> tf.Tensor:
    return x * tf.math.log(lam) - lam - tf.math.lgamma(x + 1.0)


@check_shapes(
    "x: [broadcast shape...]",
    "scale: [broadcast shape...]",
    "return: [shape...]",
)
def exponential(x: TensorType, scale: TensorType) -> tf.Tensor:
    return -x / scale - tf.math.log(scale)


@check_shapes(
    "x: [broadcast shape...]",
    "shape: [broadcast shape...]",
    "scale: [broadcast shape...]",
    "return: [shape...]",
)
def gamma(x: TensorType, shape: TensorType, scale: TensorType) -> tf.Tensor:
    return (
        -shape * tf.math.log(scale)
        - tf.math.lgamma(shape)
        + (shape - 1.0) * tf.math.log(x)
        - x / scale
    )


@check_shapes(
    "x: [broadcast shape...]",
    "mean: [broadcast shape...]",
    "scale: [broadcast shape...]",
    "df: [broadcast shape...]",
    "return: [shape...]",
)
def student_t(x: TensorType, mean: TensorType, scale: TensorType, df: TensorType) -> tf.Tensor:
    df = to_default_float(df)
    const = (
        tf.math.lgamma((df + 1.0) * 0.5)
        - tf.math.lgamma(df * 0.5)
        - 0.5 * (tf.math.log(tf.square(scale)) + tf.math.log(df) + np.log(np.pi))
    )
    return const - 0.5 * (df + 1.0) * tf.math.log(
        1.0 + (1.0 / df) * (tf.square((x - mean) / scale))
    )


@check_shapes(
    "x: [broadcast shape...]",
    "alpha: [broadcast shape...]",
    "beta: [broadcast shape...]",
    "return: [shape...]",
)
def beta(x: TensorType, alpha: TensorType, beta: TensorType) -> tf.Tensor:
    # need to clip x, since log of 0 is nan...
    x = tf.clip_by_value(x, 1e-6, 1 - 1e-6)
    return (
        (alpha - 1.0) * tf.math.log(x)
        + (beta - 1.0) * tf.math.log(1.0 - x)
        + tf.math.lgamma(alpha + beta)
        - tf.math.lgamma(alpha)
        - tf.math.lgamma(beta)
    )


@check_shapes(
    "x: [broadcast shape...]",
    "mu: [broadcast shape...]",
    "sigma: [broadcast shape...]",
    "return: [shape...]",
)
def laplace(x: TensorType, mu: TensorType, sigma: TensorType) -> tf.Tensor:
    return -tf.abs(mu - x) / sigma - tf.math.log(2.0 * sigma)


@check_shapes(
    "x: [D, broadcast N]",
    "mu: [D, broadcast N]",
    "L: [D, D]",
    "return: [N]",
)
def multivariate_normal(x: TensorType, mu: TensorType, L: TensorType) -> tf.Tensor:
    """
    Computes the log-density of a multivariate normal.

    :param x: sample(s) for which we want the density
    :param mu: mean(s) of the normal distribution
    :param L: Cholesky decomposition of the covariance matrix
    :return: log densities
    """

    d = x - mu
    alpha = tf.linalg.triangular_solve(L, d, lower=True)
    num_dims = tf.cast(tf.shape(d)[0], L.dtype)
    p = -0.5 * tf.reduce_sum(tf.square(alpha), 0)
    p -= 0.5 * num_dims * np.log(2 * np.pi)
    p -= tf.reduce_sum(tf.math.log(tf.linalg.diag_part(L)))

    return p


# gpflow/mean_functions.py

# Copyright 2016-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.
"""
Alias for (deprecated) mean_functions.py.

Use functions.py instead.
"""

from .functions import (
    Additive,
    Constant,
    Identity,
    Linear,
    MeanFunction,
    Polynomial,
    Product,
    SwitchedMeanFunction,
    Zero,
)

__all__ = [
    "Additive",
    "Constant",
    "Identity",
    "Linear",
    "MeanFunction",
    "Polynomial",
    "Product",
    "SwitchedMeanFunction",
    "Zero",
]


# gpflow/mypy_flags.py

# Copyright 2022 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.
"""
Code for printing flags for mypy, depending on library versions.
"""
from .type_flags import compute_mypy_flags


def print_mypy_flags() -> None:
    print(compute_mypy_flags())


if __name__ == "__main__":
    print_mypy_flags()


# gpflow/posteriors.py

# Copyright 2016-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.

import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Tuple, Type, Union, cast

import tensorflow as tf

from . import covariances, kernels, mean_functions
from .base import MeanAndVariance, Module, RegressionData, TensorType
from .conditionals.util import (
    base_conditional,
    base_conditional_with_lm,
    expand_independent_outputs,
    fully_correlated_conditional,
    independent_interdomain_conditional,
    mix_latent_gp,
    separate_independent_conditional_implementation,
)
from .config import default_float, default_jitter
from .covariances import Kuf, Kuu
from .experimental.check_shapes import check_shapes, inherit_check_shapes
from .inducing_variables import (
    FallbackSeparateIndependentInducingVariables,
    FallbackSharedIndependentInducingVariables,
    InducingPoints,
    InducingVariables,
    SeparateIndependentInducingVariables,
    SharedIndependentInducingVariables,
)
from .kernels import Kernel
from .likelihoods import Gaussian
from .mean_functions import MeanFunction
from .utilities import Dispatcher, add_likelihood_noise_cov, assert_params_false
from .utilities.ops import eye, leading_transpose


class _QDistribution(Module):
    """
    Base class for our parametrization of q(u) in the `AbstractPosterior`.
    Internal - do not rely on this outside of GPflow.
    """


class _DeltaDist(_QDistribution):
    @check_shapes(
        "q_mu: [M, L]",
    )
    def __init__(self, q_mu: TensorType) -> None:
        self.q_mu = q_mu

    @property
    def q_sqrt(self) -> Optional[tf.Tensor]:
        return None


class _DiagNormal(_QDistribution):
    @check_shapes(
        "q_mu: [M, L]",
        "q_sqrt: [M, L]",
    )
    def __init__(self, q_mu: TensorType, q_sqrt: TensorType) -> None:
        self.q_mu = q_mu
        self.q_sqrt = q_sqrt


class _MvNormal(_QDistribution):
    @check_shapes(
        "q_mu: [M, L]",
        "q_sqrt: [L, M, M]  # lower-triangular",
    )
    def __init__(self, q_mu: TensorType, q_sqrt: TensorType) -> None:
        self.q_mu = q_mu
        self.q_sqrt = q_sqrt


class PrecomputeCacheType(enum.Enum):
    """
    - `PrecomputeCacheType.TENSOR` (or `"tensor"`): Precomputes the cached
      quantities and stores them as tensors (which allows differentiating
      through the prediction). This is the default.
    - `PrecomputeCacheType.VARIABLE` (or `"variable"`): Precomputes the cached
      quantities and stores them as variables, which allows for updating
      their values without changing the compute graph (relevant for AOT
      compilation).
    - `PrecomputeCacheType.NOCACHE` (or `"nocache"` or `None`): Avoids
      immediate cache computation. This is useful for avoiding extraneous
      computations when you only want to call the posterior's
      `fused_predict_f` method.
    """

    TENSOR = "tensor"
    VARIABLE = "variable"
    NOCACHE = "nocache"


@dataclass
class PrecomputedValue:
    value: tf.Tensor
    """
    The precomputed value itself.
    """

    axis_dynamic: Tuple[bool, ...]
    """
    A tuple with one element per dimension of `value`. That element is `True` if that dimension
    of `value` might change size.
    """

    def __post_init__(self) -> None:
        tf.debugging.assert_rank(
            self.value,
            len(self.axis_dynamic),
            "axis_dynamic must have one element per dimension of value.",
        )

    @staticmethod
    @check_shapes(
        "alpha: [M_L_or_L_M_M...]",
        "Qinv: [M_M_or_L_M_M...]",
    )
    def wrap_alpha_Qinv(alpha: TensorType, Qinv: TensorType) -> Tuple["PrecomputedValue", ...]:
        """
        Wraps `alpha` and `Qinv` in `PrecomputedValue`\ s.
        """
        one_dynamic = False
        L_dynamic = False
        M_dynamic = False  # TODO(jesper): Support variable number of inducing points?

        alpha_rank = tf.rank(alpha)
        if alpha_rank == 2:
            alpha_dynamic: Tuple[bool, ...] = (M_dynamic, L_dynamic)
        elif alpha_rank == 3:
            alpha_dynamic = (L_dynamic, M_dynamic, one_dynamic)
        else:
            raise AssertionError(f"Unknown rank of alpha {alpha_rank}.")

        Qinv_rank = tf.rank(Qinv)
        if Qinv_rank == 2:
            Qinv_dynamic: Tuple[bool, ...] = (M_dynamic, M_dynamic)
        elif Qinv_rank == 3:
            Qinv_dynamic = (L_dynamic, M_dynamic, M_dynamic)
        else:
            raise AssertionError(f"Unknown rank of Qinv {Qinv_rank}.")

        return (
            PrecomputedValue(alpha, alpha_dynamic),
            PrecomputedValue(Qinv, Qinv_dynamic),
        )


def _validate_precompute_cache_type(
    value: Union[None, PrecomputeCacheType, str]
) -> PrecomputeCacheType:
    if value is None:
        return PrecomputeCacheType.NOCACHE
    elif isinstance(value, PrecomputeCacheType):
        return value
    elif isinstance(value, str):
        return PrecomputeCacheType(value.lower())
    else:
        raise ValueError(
            f"{value} is not a valid PrecomputeCacheType."
            " Valid options: 'tensor', 'variable', 'nocache' (or None)."
        )


class AbstractPosterior(Module, ABC):
    @check_shapes(
        "X_data: [N_D_or_M_D_P...]",
    )
    def __init__(
        self,
        kernel: Kernel,
        X_data: Union[tf.Tensor, InducingVariables],
        cache: Optional[Tuple[tf.Tensor, ...]] = None,
        mean_function: Optional[mean_functions.MeanFunction] = None,
    ) -> None:
        """
        Users should use `create_posterior` to create instances of concrete
        subclasses of this AbstractPosterior class instead of calling this
        constructor directly. For `create_posterior` to be able to correctly
        instantiate subclasses, developers need to ensure their subclasses
        don't change the constructor signature.
        """
        super().__init__()

        self.kernel = kernel
        self.X_data = X_data
        self.cache = cache
        self.mean_function = mean_function

        self._precompute_cache: Optional[PrecomputeCacheType] = None

    @check_shapes(
        "Xnew: [batch..., D]",
        "mean: [batch..., Q]",
        "return: [batch..., Q]",
    )
    def _add_mean_function(self, Xnew: TensorType, mean: TensorType) -> tf.Tensor:
        if self.mean_function is None:
            return mean
        else:
            return mean + self.mean_function(Xnew)

    @abstractmethod
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        """
        Precompute a cache.

        The result of this method will later be passed to `_conditional_with_precompute` as the
        `cache` argument.
        """

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def fused_predict_f(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, including mean_function
        Does not make use of caching
        """
        mean, cov = self._conditional_fused(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )
        return self._add_mean_function(Xnew, mean), cov

    @abstractmethod
    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function
        Does not make use of caching
        """

    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def predict_f(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, including mean_function.
        Relies on precomputed alpha and Qinv (see _precompute method)
        """
        if self.cache is None:
            raise ValueError(
                "Cache has not been precomputed yet. Call update_cache first or use fused_predict_f"
            )
        mean, cov = self._conditional_with_precompute(
            self.cache, Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )
        return self._add_mean_function(Xnew, mean), cov

    @abstractmethod
    @check_shapes(
        "Xnew: [batch..., N, D]",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """

    def update_cache(self, precompute_cache: Optional[PrecomputeCacheType] = None) -> None:
        """
        Sets the cache depending on the value of `precompute_cache` to a
        `tf.Tensor`, `tf.Variable`, or clears the cache. If `precompute_cache`
        is not given, the setting defaults to the most-recently-used one.
        """
        if precompute_cache is None:
            if self._precompute_cache is None:
                raise ValueError(
                    "You must pass precompute_cache explicitly"
                    " (the cache had not been updated before)."
                )
            precompute_cache = self._precompute_cache
        else:
            self._precompute_cache = precompute_cache

        if precompute_cache is PrecomputeCacheType.NOCACHE:
            self.cache = None

        elif precompute_cache is PrecomputeCacheType.TENSOR:
            self.cache = tuple(c.value for c in self._precompute())

        elif precompute_cache is PrecomputeCacheType.VARIABLE:
            cache = self._precompute()

            if self.cache is not None and all(isinstance(c, tf.Variable) for c in self.cache):
                # re-use existing variables
                for cache_var, c in zip(self.cache, cache):
                    cache_var.assign(c.value)
            else:  # create variables
                shapes = [
                    [None if d else s for d, s in zip(c.axis_dynamic, tf.shape(c.value))]
                    for c in cache
                ]
                self.cache = tuple(
                    tf.Variable(c.value, trainable=False, shape=s) for c, s in zip(cache, shapes)
                )


class GPRPosterior(AbstractPosterior):
    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, Q]",
    )
    def __init__(
        self,
        kernel: Kernel,
        data: RegressionData,
        likelihood: Gaussian,
        mean_function: MeanFunction,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ) -> None:
        X, Y = data
        super().__init__(kernel, X, mean_function=mean_function)
        self.Y_data = Y
        self.likelihood = likelihood

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """
        assert_params_false(self._conditional_with_precompute, full_output_cov=full_output_cov)

        (alpha,) = cache
        (Qinv,) = cache

        Kmn = self.kernel(self.X_data, Xnew)
        # compute kernel stuff
        num_func = tf.shape(self.Y_data)[-1]  # R
        N = tf.shape(Kmn)[-1]

        # get the leading dims in Kmn to the front of the tensor Kmn
        K = tf.rank(Kmn)
        perm = tf.concat(
            [
                tf.reshape(tf.range(1, K - 1), [K - 2]),  # leading dims (...)
                tf.reshape(0, [1]),  # [M]
                tf.reshape(K - 1, [1]),
            ],
            0,
        )  # [N]
        Kmn = tf.transpose(Kmn, perm)  # [..., M, N]
        leading_dims = tf.shape(Kmn)[:-2]

        # get the leading dims in Knm to the front of the tensor Knm
        Knm = leading_transpose(Kmn, [..., -1, -2])

        assert self.mean_function is not None
        Knn = self.kernel(Xnew, full_cov=full_cov)
        err = self.Y_data - self.mean_function(self.X_data)

        mean = Knm @ alpha @ err

        # The GPR model only has a single latent GP.
        if full_cov:
            cov = Knn - Knm @ Qinv @ Kmn  # [..., N, N]
            cov_shape = tf.concat([leading_dims, [num_func, N, N]], 0)
            cov = tf.broadcast_to(tf.expand_dims(cov, -3), cov_shape)  # [..., R, N, N]

        else:
            cov = Knn - tf.einsum("...ij,...ji->...i", Knm @ Qinv, Kmn)  # [..., N]
            cov_shape = tf.concat([leading_dims, [num_func, N]], 0)  # [..., R, N]
            cov = tf.broadcast_to(tf.expand_dims(cov, -2), cov_shape)  # [..., R, N]
            cov = tf.linalg.adjoint(cov)

        return mean, cov

    @check_shapes(
        "return[0].value: [M, M]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        X_data = cast(tf.Tensor, self.X_data)
        Kmm = self.kernel(X_data)
        Kmm_plus_s = add_likelihood_noise_cov(Kmm, self.likelihood, X_data)

        Lm = tf.linalg.cholesky(Kmm_plus_s)
        Kmm_plus_s_inv = tf.linalg.cholesky_solve(Lm, tf.eye(tf.shape(X_data)[0], dtype=Lm.dtype))

        M = X_data.shape[0]
        M_dynamic = M is None

        return (PrecomputedValue(Kmm_plus_s_inv, (M_dynamic, M_dynamic)),)

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function
        Does not make use of caching
        """
        assert_params_false(self._conditional_fused, full_output_cov=full_output_cov)

        # taken directly from the deprecated GPR implementation
        assert self.mean_function is not None
        err = self.Y_data - self.mean_function(self.X_data)

        Kmm = self.kernel(self.X_data)
        Knn = self.kernel(Xnew, full_cov=full_cov)
        Kmn = self.kernel(self.X_data, Xnew)
        Kmm_plus_s = add_likelihood_noise_cov(Kmm, self.likelihood, self.X_data)

        return base_conditional(
            Kmn, Kmm_plus_s, Knn, err, full_cov=full_cov, white=False
        )  # [N, P], [N, P] or [P, N, N]


class SGPRPosterior(AbstractPosterior):
    """
    This class represents posteriors which can be derived from SGPR
    models to compute faster predictions on unseen points.
    """

    @check_shapes(
        "data[0]: [N, D]",
        "data[1]: [N, Q]",
        "inducing_variable: [M, D, 1]",
    )
    def __init__(
        self,
        kernel: Kernel,
        data: RegressionData,
        inducing_variable: InducingPoints,
        likelihood: Gaussian,
        num_latent_gps: int,
        mean_function: MeanFunction,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ) -> None:
        X, Y = data
        super().__init__(kernel, X, mean_function=mean_function)
        self.Y_data = Y
        self.likelihood = likelihood
        self.inducing_variable = inducing_variable
        self.num_latent_gps = num_latent_gps

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @check_shapes(
        "return[0].value: [M, L]",
        "return[1].value: [M, M]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        # taken directly from the deprecated SGPR implementation
        num_inducing = self.inducing_variable.num_inducing
        assert self.mean_function is not None
        err = self.Y_data - self.mean_function(self.X_data)
        kuf = Kuf(self.inducing_variable, self.kernel, self.X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())

        sigma_sq = tf.squeeze(self.likelihood.variance_at(self.X_data), axis=-1)
        sigma = tf.sqrt(sigma_sq)

        L = tf.linalg.cholesky(kuu)  # cache alpha, qinv
        A = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        B = tf.linalg.matmul(A, A, transpose_b=True) + tf.eye(
            num_inducing, dtype=default_float()
        )  # cache qinv
        LB = tf.linalg.cholesky(B)  # cache alpha
        Aerr = tf.linalg.matmul(A, err / sigma[..., None])
        c = tf.linalg.triangular_solve(LB, Aerr, lower=True)  # cache alpha

        # get intermediate variables
        Linv = tf.linalg.triangular_solve(L, tf.eye(num_inducing, dtype=default_float()))
        LBinv = tf.linalg.triangular_solve(LB, tf.eye(num_inducing, dtype=default_float()))
        Binv = tf.linalg.inv(B)  # naive...can do better?
        tmp = tf.eye(num_inducing, dtype=default_float()) - Binv

        # calculate cached values
        LinvT = tf.transpose(Linv)
        alpha = LinvT @ tf.transpose(LBinv) @ c
        Qinv = LinvT @ tmp @ Linv

        return PrecomputedValue.wrap_alpha_Qinv(alpha, Qinv)

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        """
        Computes predictive mean and (co)variance at Xnew, *excluding* mean_function.
        Relies on cached alpha and Qinv.
        """
        assert_params_false(self._conditional_with_precompute, full_output_cov=full_output_cov)

        alpha, Qinv = cache

        Kus = Kuf(self.inducing_variable, self.kernel, Xnew)
        Knn = self.kernel(Xnew, full_cov=full_cov)

        Ksu = tf.transpose(Kus)
        mean = Ksu @ alpha

        if full_cov:
            var = Knn - Ksu @ Qinv @ Kus
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            Kfu_Qinv_Kuf = tf.reduce_sum(Kus * tf.matmul(Qinv, Kus), axis=-2)
            var = Knn - Kfu_Qinv_Kuf
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        return mean, var

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        """
        Compute the mean and variance of the latent function at some new points
        Xnew. Does not make use of caching
        """
        # taken directly from the deprecated SGPR implementation

        assert_params_false(self._conditional_fused, full_output_cov=full_output_cov)

        num_inducing = self.inducing_variable.num_inducing
        assert self.mean_function is not None
        err = self.Y_data - self.mean_function(self.X_data)
        kuf = Kuf(self.inducing_variable, self.kernel, self.X_data)
        kuu = Kuu(self.inducing_variable, self.kernel, jitter=default_jitter())
        Kus = Kuf(self.inducing_variable, self.kernel, Xnew)

        sigma_sq = tf.squeeze(self.likelihood.variance_at(self.X_data), axis=-1)
        sigma = tf.sqrt(sigma_sq)

        L = tf.linalg.cholesky(kuu)  # cache alpha, qinv
        A = tf.linalg.triangular_solve(L, kuf / sigma, lower=True)
        B = tf.linalg.matmul(A, A, transpose_b=True) + tf.eye(
            num_inducing, dtype=default_float()
        )  # cache qinv
        LB = tf.linalg.cholesky(B)  # cache alpha
        Aerr = tf.linalg.matmul(A, err / sigma[..., None])
        c = tf.linalg.triangular_solve(LB, Aerr, lower=True)  # cache alpha
        tmp1 = tf.linalg.triangular_solve(L, Kus, lower=True)
        tmp2 = tf.linalg.triangular_solve(LB, tmp1, lower=True)
        mean = tf.linalg.matmul(tmp2, c, transpose_a=True)
        if full_cov:
            var = (
                self.kernel(Xnew)
                + tf.linalg.matmul(tmp2, tmp2, transpose_a=True)
                - tf.linalg.matmul(tmp1, tmp1, transpose_a=True)
            )
            var = tf.tile(var[None, ...], [self.num_latent_gps, 1, 1])  # [P, N, N]
        else:
            var = (
                self.kernel(Xnew, full_cov=False)
                + tf.reduce_sum(tf.square(tmp2), 0)
                - tf.reduce_sum(tf.square(tmp1), 0)
            )
            var = tf.tile(var[:, None], [1, self.num_latent_gps])

        return mean, var


class VGPPosterior(AbstractPosterior):
    @check_shapes(
        "X: [N, D]",
        "q_mu: [N, P]",
        "q_sqrt: [N_P_or_P_N_N...]",
    )
    def __init__(
        self,
        kernel: Kernel,
        X: tf.Tensor,
        q_mu: tf.Tensor,
        q_sqrt: tf.Tensor,
        mean_function: Optional[mean_functions.MeanFunction] = None,
        white: bool = True,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ) -> None:
        super().__init__(kernel, X, mean_function=mean_function)
        self.q_mu = q_mu
        self.q_sqrt = q_sqrt
        self.white = white

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        assert_params_false(self._conditional_with_precompute, full_output_cov=full_output_cov)

        (Lm,) = cache
        Kmn = self.kernel(self.X_data, Xnew)  # [M, ..., N]
        Knn = self.kernel(
            Xnew, full_cov=full_cov
        )  # [..., N] (full_cov = False) or [..., N, N] (True)

        return base_conditional_with_lm(
            Kmn=Kmn,
            Lm=Lm,
            Knn=Knn,
            f=self.q_mu,
            full_cov=full_cov,
            q_sqrt=self.q_sqrt,
            white=self.white,
        )

    @check_shapes(
        "return[0].value: [M, M]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        X_data = cast(tf.Tensor, self.X_data)
        Kmm = self.kernel(X_data) + eye(
            tf.shape(X_data)[-2], value=default_jitter(), dtype=X_data.dtype
        )  # [..., M, M]
        Lm = tf.linalg.cholesky(Kmm)

        M = X_data.shape[0]
        M_dynamic = M is None

        return (PrecomputedValue(Lm, (M_dynamic, M_dynamic)),)

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        temp_cache = tuple(c.value for c in self._precompute())
        return self._conditional_with_precompute(temp_cache, Xnew, full_cov, full_output_cov)


class BasePosterior(AbstractPosterior):
    @check_shapes(
        "inducing_variable: [M, D, broadcast P]",
        "q_mu: [N, P]",
        "q_sqrt: [N_P_or_P_N_N...]",
    )
    def __init__(
        self,
        kernel: Kernel,
        inducing_variable: InducingVariables,
        q_mu: tf.Tensor,
        q_sqrt: tf.Tensor,
        whiten: bool = True,
        mean_function: Optional[mean_functions.MeanFunction] = None,
        *,
        precompute_cache: Optional[PrecomputeCacheType],
    ):

        super().__init__(kernel, inducing_variable, mean_function=mean_function)
        self.whiten = whiten
        self._set_qdist(q_mu, q_sqrt)

        if precompute_cache is not None:
            self.update_cache(precompute_cache)

    @property  # type: ignore[misc]
    @check_shapes(
        "return: [N, P]",
    )
    def q_mu(self) -> tf.Tensor:
        return self._q_dist.q_mu

    @property  # type: ignore[misc]
    @check_shapes(
        "return: [N_P_or_P_N_N...]",
    )
    def q_sqrt(self) -> tf.Tensor:
        return self._q_dist.q_sqrt

    @check_shapes(
        "q_mu: [N, P]",
        "q_sqrt: [N_P_or_P_N_N...]",
    )
    def _set_qdist(self, q_mu: TensorType, q_sqrt: TensorType) -> None:
        if q_sqrt is None:
            self._q_dist = _DeltaDist(q_mu)
        elif len(q_sqrt.shape) == 2:  # q_diag
            self._q_dist = _DiagNormal(q_mu, q_sqrt)
        else:
            self._q_dist = _MvNormal(q_mu, q_sqrt)

    @check_shapes(
        "return[0].value: [M_L_or_L_M_M...]",
        "return[1].value: [L, M, M]",
    )
    def _precompute(self) -> Tuple[PrecomputedValue, ...]:
        Kuu = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [(R), M, M]
        q_mu = self._q_dist.q_mu

        if Kuu.shape.ndims == 4:
            ML = tf.reduce_prod(tf.shape(Kuu)[:2])
            Kuu = tf.reshape(Kuu, [ML, ML])
        if Kuu.shape.ndims == 3:
            q_mu = tf.linalg.adjoint(self._q_dist.q_mu)[..., None]  # [..., R, M, 1]
        L = tf.linalg.cholesky(Kuu)

        if not self.whiten:
            # alpha = Kuu⁻¹ q_mu
            alpha = tf.linalg.cholesky_solve(L, q_mu)
        else:
            # alpha = L⁻ᵀ q_mu
            alpha = tf.linalg.triangular_solve(L, q_mu, adjoint=True)
        # predictive mean = Kfu alpha
        # predictive variance = Kff - Kfu Qinv Kuf
        # S = q_sqrt q_sqrtᵀ
        I = tf.eye(tf.shape(L)[-1], dtype=L.dtype)
        if isinstance(self._q_dist, _DeltaDist):
            B = I
        else:
            if not self.whiten:
                # Qinv = Kuu⁻¹ - Kuu⁻¹ S Kuu⁻¹
                #      = Kuu⁻¹ - L⁻ᵀ L⁻¹ S L⁻ᵀ L⁻¹
                #      = L⁻ᵀ (I - L⁻¹ S L⁻ᵀ) L⁻¹
                #      = L⁻ᵀ B L⁻¹
                if isinstance(self._q_dist, _DiagNormal):
                    q_sqrt = tf.linalg.diag(tf.linalg.adjoint(self._q_dist.q_sqrt))
                elif isinstance(self._q_dist, _MvNormal):
                    q_sqrt = self._q_dist.q_sqrt
                Linv_qsqrt = tf.linalg.triangular_solve(L, q_sqrt)
                Linv_cov_u_LinvT = tf.matmul(Linv_qsqrt, Linv_qsqrt, transpose_b=True)
            else:
                if isinstance(self._q_dist, _DiagNormal):
                    Linv_cov_u_LinvT = tf.linalg.diag(tf.linalg.adjoint(self._q_dist.q_sqrt ** 2))
                elif isinstance(self._q_dist, _MvNormal):
                    q_sqrt = self._q_dist.q_sqrt
                    Linv_cov_u_LinvT = tf.matmul(q_sqrt, q_sqrt, transpose_b=True)
                # Qinv = Kuu⁻¹ - L⁻ᵀ S L⁻¹
                # Linv = (L⁻¹ I) = solve(L, I)
                # Kinv = Linvᵀ @ Linv
            B = I - Linv_cov_u_LinvT
        LinvT_B = tf.linalg.triangular_solve(L, B, adjoint=True)
        B_Linv = tf.linalg.adjoint(LinvT_B)
        Qinv = tf.linalg.triangular_solve(L, B_Linv, adjoint=True)

        M, L = tf.unstack(tf.shape(self._q_dist.q_mu), num=2)
        Qinv = tf.broadcast_to(Qinv, [L, M, M])

        return PrecomputedValue.wrap_alpha_Qinv(alpha, Qinv)


class IndependentPosterior(BasePosterior):
    @check_shapes(
        "mean: [batch..., N, P]",
        "cov: [batch..., P, N, N] if full_cov",
        "cov: [batch..., N, P] if not full_cov",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _post_process_mean_and_cov(
        self, mean: TensorType, cov: TensorType, full_cov: bool, full_output_cov: bool
    ) -> MeanAndVariance:
        return mean, expand_independent_outputs(cov, full_cov, full_output_cov)

    @check_shapes(
        "Xnew: [N, D]",
        "return: [broadcast P, N, N] if full_cov",
        "return: [broadcast P, N] if (not full_cov)",
    )
    def _get_Kff(self, Xnew: TensorType, full_cov: bool) -> tf.Tensor:

        # TODO: this assumes that Xnew has shape [N, D] and no leading dims

        if isinstance(self.kernel, (kernels.SeparateIndependent, kernels.IndependentLatent)):
            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would
            # return
            # if full_cov: [P, N, N] -- this is what we want
            # else: [N, P] instead of [P, N] as we get from the explicit stack below
            Kff = tf.stack([k(Xnew, full_cov=full_cov) for k in self.kernel.kernels], axis=0)
        elif isinstance(self.kernel, kernels.MultioutputKernel):
            # effectively, SharedIndependent path
            Kff = self.kernel.kernel(Xnew, full_cov=full_cov)
            # NOTE calling kernel(Xnew, full_cov=full_cov, full_output_cov=False) directly would
            # return
            # if full_cov: [P, N, N] instead of [N, N]
            # else: [N, P] instead of [N]
        else:
            # standard ("single-output") kernels
            Kff = self.kernel(Xnew, full_cov=full_cov)  # [N, N] if full_cov else [N]

        return Kff

    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        # Qinv: [L, M, M]
        # alpha: [M, L]
        alpha, Qinv = cache

        Kuf = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [(R), M, N]
        Kff = self._get_Kff(Xnew, full_cov)

        mean = tf.matmul(Kuf, alpha, transpose_a=True)
        if Kuf.shape.ndims == 3:
            mean = tf.linalg.adjoint(tf.squeeze(mean, axis=-1))

        if full_cov:
            Kfu_Qinv_Kuf = tf.matmul(Kuf, Qinv @ Kuf, transpose_a=True)
            cov = Kff - Kfu_Qinv_Kuf
        else:
            # [Aᵀ B]_ij = Aᵀ_ik B_kj = A_ki B_kj
            # TODO check whether einsum is faster now?
            Kfu_Qinv_Kuf = tf.reduce_sum(Kuf * tf.matmul(Qinv, Kuf), axis=-2)
            cov = Kff - Kfu_Qinv_Kuf
            cov = tf.linalg.adjoint(cov)

        return self._post_process_mean_and_cov(mean, cov, full_cov, full_output_cov)


class IndependentPosteriorSingleOutput(IndependentPosterior):
    # could almost be the same as IndependentPosteriorMultiOutput ...
    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        # same as IndependentPosteriorMultiOutput, Shared~/Shared~ branch, except for following
        # line:
        Knn = self.kernel(Xnew, full_cov=full_cov)

        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, M]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, N]

        fmean, fvar = base_conditional(
            Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
        )  # [N, P],  [P, N, N] or [N, P]
        return self._post_process_mean_and_cov(fmean, fvar, full_cov, full_output_cov)


class IndependentPosteriorMultiOutput(IndependentPosterior):
    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        if isinstance(self.X_data, SharedIndependentInducingVariables) and isinstance(
            self.kernel, kernels.SharedIndependent
        ):
            # same as IndependentPosteriorSingleOutput except for following line
            Knn = self.kernel.kernel(Xnew, full_cov=full_cov)
            # we don't call self.kernel() directly as that would do unnecessary tiling

            Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, M]
            Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, N]

            fmean, fvar = base_conditional(
                Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
            )  # [N, P],  [P, N, N] or [N, P]
        else:
            # this is the messy thing with tf.map_fn, cleaned up by the
            # st/clean_up_broadcasting_conditionals branch

            # Following are: [P, M, M]  -  [P, M, N]  -  [P, N](x N)
            Kmms = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [P, M, M]
            Kmns = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [P, M, N]
            if isinstance(self.kernel, kernels.Combination):
                kernel_list = self.kernel.kernels
            else:
                kernel_list = [self.kernel.kernel] * len(self.X_data.inducing_variable_list)
            Knns = tf.stack(
                [k.K(Xnew) if full_cov else k.K_diag(Xnew) for k in kernel_list], axis=0
            )

            fmean, fvar = separate_independent_conditional_implementation(
                Kmns,
                Kmms,
                Knns,
                self.q_mu,
                q_sqrt=self.q_sqrt,
                full_cov=full_cov,
                white=self.whiten,
            )

        return self._post_process_mean_and_cov(fmean, fvar, full_cov, full_output_cov)


class LinearCoregionalizationPosterior(IndependentPosteriorMultiOutput):
    @check_shapes(
        "mean: [batch..., N, L]",
        "cov: [batch..., L, N, N] if full_cov",
        "cov: [batch..., N, L] if not full_cov",
        "return[0]: [batch..., N, P]",
        "return[1]: [batch..., N, P, N, P] if full_cov and full_output_cov",
        "return[1]: [batch..., N, P, P] if (not full_cov) and full_output_cov",
        "return[1]: [batch..., P, N, N] if full_cov and (not full_output_cov)",
        "return[1]: [batch..., N, P] if (not full_cov) and (not full_output_cov)",
    )
    def _post_process_mean_and_cov(
        self, mean: TensorType, cov: TensorType, full_cov: bool, full_output_cov: bool
    ) -> MeanAndVariance:
        cov = expand_independent_outputs(cov, full_cov, full_output_cov=False)
        mean, cov = mix_latent_gp(self.kernel.W, mean, cov, full_cov, full_output_cov)
        return mean, cov


class FullyCorrelatedPosterior(BasePosterior):
    @inherit_check_shapes
    def _conditional_with_precompute(
        self,
        cache: Tuple[tf.Tensor, ...],
        Xnew: TensorType,
        full_cov: bool = False,
        full_output_cov: bool = False,
    ) -> MeanAndVariance:
        # TODO: this assumes that Xnew has shape [N, D] and no leading dims

        # Qinv: [L, M, M]
        # alpha: [M, L]
        alpha, Qinv = cache

        Kuf = covariances.Kuf(self.X_data, self.kernel, Xnew)
        assert Kuf.shape.ndims == 4
        M, L, N, K = tf.unstack(tf.shape(Kuf), num=Kuf.shape.ndims, axis=0)
        Kuf = tf.reshape(Kuf, (M * L, N * K))

        kernel: kernels.MultioutputKernel = self.kernel
        Kff = kernel(Xnew, full_cov=full_cov, full_output_cov=full_output_cov)
        # full_cov=True and full_output_cov=True: [N, P, N, P]
        # full_cov=True and full_output_cov=False: [P, N, N]
        # full_cov=False and full_output_cov=True: [N, P, P]
        # full_cov=False and full_output_cov=False: [N, P]
        if full_cov == full_output_cov:
            new_shape = (N * K, N * K) if full_cov else (N * K,)
            Kff = tf.reshape(Kff, new_shape)

        N = tf.shape(Xnew)[0]
        K = tf.shape(Kuf)[-1] // N

        mean = tf.matmul(Kuf, alpha, transpose_a=True)
        if Kuf.shape.ndims == 3:
            mean = tf.linalg.adjoint(tf.squeeze(mean, axis=-1))

        if not full_cov and not full_output_cov:
            # fully diagonal case in both inputs and outputs
            # [Aᵀ B]_ij = Aᵀ_ik B_kj = A_ki B_kj
            # TODO check whether einsum is faster now?
            Kfu_Qinv_Kuf = tf.reduce_sum(Kuf * tf.matmul(Qinv, Kuf), axis=-2)
        else:
            Kfu_Qinv_Kuf = tf.matmul(Kuf, Qinv @ Kuf, transpose_a=True)
            if not (full_cov and full_output_cov):
                # diagonal in either inputs or outputs
                new_shape = tf.concat([tf.shape(Kfu_Qinv_Kuf)[:-2], (N, K, N, K)], axis=0)
                Kfu_Qinv_Kuf = tf.reshape(Kfu_Qinv_Kuf, new_shape)
                if full_cov:
                    # diagonal in outputs: move outputs to end
                    tmp = tf.linalg.diag_part(tf.einsum("...ijkl->...ikjl", Kfu_Qinv_Kuf))
                elif full_output_cov:
                    # diagonal in inputs: move inputs to end
                    tmp = tf.linalg.diag_part(tf.einsum("...ijkl->...jlik", Kfu_Qinv_Kuf))
                Kfu_Qinv_Kuf = tf.einsum("...ijk->...kij", tmp)  # move diagonal dim to [-3]
        cov = Kff - Kfu_Qinv_Kuf

        if not full_cov and not full_output_cov:
            cov = tf.linalg.adjoint(cov)

        mean = tf.reshape(mean, (N, K))
        if full_cov == full_output_cov:
            cov_shape = (N, K, N, K) if full_cov else (N, K)
        else:
            cov_shape = (K, N, N) if full_cov else (N, K, K)
        cov = tf.reshape(cov, cov_shape)

        return mean, cov

    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [M, L, M, L]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, L, N, P]
        kernel: kernels.MultioutputKernel = self.kernel
        Knn = kernel(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )  # [N, P](x N)x P  or  [N, P](x P)

        M, L, N, K = tf.unstack(tf.shape(Kmn), num=Kmn.shape.ndims, axis=0)
        Kmm = tf.reshape(Kmm, (M * L, M * L))

        if full_cov == full_output_cov:
            Kmn = tf.reshape(Kmn, (M * L, N * K))
            Knn = tf.reshape(Knn, (N * K, N * K)) if full_cov else tf.reshape(Knn, (N * K,))
            mean, cov = base_conditional(
                Kmn, Kmm, Knn, self.q_mu, full_cov=full_cov, q_sqrt=self.q_sqrt, white=self.whiten
            )  # [K, 1], [1, K](x NK)
            mean = tf.reshape(mean, (N, K))
            cov = tf.reshape(cov, (N, K, N, K) if full_cov else (N, K))
        else:
            Kmn = tf.reshape(Kmn, (M * L, N, K))
            mean, cov = fully_correlated_conditional(
                Kmn,
                Kmm,
                Knn,
                self.q_mu,
                full_cov=full_cov,
                full_output_cov=full_output_cov,
                q_sqrt=self.q_sqrt,
                white=self.whiten,
            )
        return mean, cov


class FallbackIndependentLatentPosterior(FullyCorrelatedPosterior):  # XXX
    @inherit_check_shapes
    def _conditional_fused(
        self, Xnew: TensorType, full_cov: bool = False, full_output_cov: bool = False
    ) -> MeanAndVariance:
        Kmm = covariances.Kuu(self.X_data, self.kernel, jitter=default_jitter())  # [L, M, M]
        Kmn = covariances.Kuf(self.X_data, self.kernel, Xnew)  # [M, L, N, P]
        kernel: kernels.IndependentLatent = self.kernel
        Knn = kernel(
            Xnew, full_cov=full_cov, full_output_cov=full_output_cov
        )  # [N, P](x N)x P  or  [N, P](x P)

        return independent_interdomain_conditional(
            Kmn,
            Kmm,
            Knn,
            self.q_mu,
            full_cov=full_cov,
            full_output_cov=full_output_cov,
            q_sqrt=self.q_sqrt,
            white=self.whiten,
        )


get_posterior_class = Dispatcher("get_posterior_class")


@get_posterior_class.register(kernels.Kernel, InducingVariables)
def _get_posterior_base_case(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    # independent single output
    return IndependentPosteriorSingleOutput


@get_posterior_class.register(kernels.MultioutputKernel, InducingPoints)
def _get_posterior_fully_correlated_mo(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    return FullyCorrelatedPosterior


@get_posterior_class.register(
    (kernels.SharedIndependent, kernels.SeparateIndependent),
    (SeparateIndependentInducingVariables, SharedIndependentInducingVariables),
)
def _get_posterior_independent_mo(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    # independent multi-output
    return IndependentPosteriorMultiOutput


@get_posterior_class.register(
    kernels.IndependentLatent,
    (FallbackSeparateIndependentInducingVariables, FallbackSharedIndependentInducingVariables),
)
def _get_posterior_independentlatent_mo_fallback(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    return FallbackIndependentLatentPosterior


@get_posterior_class.register(
    kernels.LinearCoregionalization,
    (SeparateIndependentInducingVariables, SharedIndependentInducingVariables),
)
def _get_posterior_linearcoregionalization_mo_efficient(
    kernel: Kernel, inducing_variable: InducingVariables
) -> Type[BasePosterior]:
    # Linear mixing---efficient multi-output
    return LinearCoregionalizationPosterior


def create_posterior(
    kernel: Kernel,
    inducing_variable: InducingVariables,
    q_mu: TensorType,
    q_sqrt: TensorType,
    whiten: bool,
    mean_function: Optional[MeanFunction] = None,
    precompute_cache: Union[PrecomputeCacheType, str, None] = PrecomputeCacheType.TENSOR,
) -> BasePosterior:
    posterior_class = get_posterior_class(kernel, inducing_variable)
    precompute_cache = _validate_precompute_cache_type(precompute_cache)
    return posterior_class(  # type: ignore[no-any-return]
        kernel,
        inducing_variable,
        q_mu,
        q_sqrt,
        whiten,
        mean_function,
        precompute_cache=precompute_cache,
    )


# gpflow/probability_distributions.py

# 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.


# Eventually, it would be nice to not have to have our own classes for
# probability distributions. The TensorFlow "distributions" framework would
# be a good replacement.
from abc import ABC, abstractmethod

from .base import TensorType
from .experimental.check_shapes import ErrorContext, Shape, check_shapes, register_get_shape


class ProbabilityDistribution(ABC):
    """
    This is the base class for a probability distributions,
    over which we take the expectations in the expectations framework.
    """

    @property
    @abstractmethod
    def shape(self) -> Shape:
        """
        Return the shape of this distribution.

        Shape should be some variation of ``[N, D]``, where:

        * ``N`` is the number of data points.
        * ``D`` is the number of input dimensions.
        """


@register_get_shape(ProbabilityDistribution)
def get_probability_distribution_shape(
    shaped: ProbabilityDistribution, context: ErrorContext
) -> Shape:
    return shaped.shape


class Gaussian(ProbabilityDistribution):
    @check_shapes(
        "mu: [N, D]",
        "cov: [N, D, D]",
    )
    def __init__(self, mu: TensorType, cov: TensorType):
        self.mu = mu
        self.cov = cov

    @property
    def shape(self) -> Shape:
        return self.mu.shape  # type: ignore[no-any-return]


class DiagonalGaussian(ProbabilityDistribution):
    @check_shapes(
        "mu: [N, D]",
        "cov: [N, D]",
    )
    def __init__(self, mu: TensorType, cov: TensorType):
        self.mu = mu
        self.cov = cov

    @property
    def shape(self) -> Shape:
        return self.mu.shape  # type: ignore[no-any-return]


class MarkovGaussian(ProbabilityDistribution):
    """
    Gaussian distribution with Markov structure.
    Only covariances and covariances between t and t+1 need to be
    parameterised. We use the solution proposed by Carl Rasmussen, i.e. to
    represent
    Var[x_t] = cov[x_t, :, :] * cov[x_t, :, :].T
    Cov[x_t, x_{t+1}] = cov[t, :, :] * cov[t+1, :, :]
    """

    @check_shapes(
        "mu: [N_plus_1, D]",
        "cov: [2, N_plus_1, D, D]",
    )
    def __init__(self, mu: TensorType, cov: TensorType):
        self.mu = mu
        self.cov = cov

    @property
    def shape(self) -> Shape:
        shape = self.mu.shape
        if shape is None:
            return shape
        N_plus_1, D = shape
        N = N_plus_1 - 1
        return N, D


# gpflow/type_flags.py

# Copyright 2022 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.
"""
Code for setting type hints, depending on library versions.
"""
import sys

from numpy import __version__ as np_version
from packaging.version import Version

NP_VERSION = Version(np_version)

NP_TYPE_CHECKING = False
"""
Whether to type-check numpy arrays at all. Defaults to False, because we don't know which
versions a client might use.
"""

GENERIC_NP_ARRAYS = (sys.version_info >= (3, 9)) and (NP_VERSION >= Version("1.22.0"))
"""
Whether to use generic numpy arrays. This is not applied at all, if type checking and not
`NP_TYPE_CHECKING`.
"""


def compute_mypy_flags() -> str:  # pragma: no cover
    from mypy.version import __version__ as mypy_version

    MYPY_VERSION = Version(mypy_version)

    flags = []

    def set_always(variable: str, value: bool) -> None:
        if value:
            flags.append("--always-true")
        else:
            flags.append("--always-false")
        flags.append(variable)

    set_always("NP_TYPE_CHECKING", True)
    set_always("GENERIC_NP_ARRAYS", NP_VERSION >= Version("1.21.0"))
    if MYPY_VERSION >= Version("0.940"):
        flags.extend(["--enable-error-code", "ignore-without-code"])

    return " ".join(flags)


# gpflow/versions.py

import pkg_resources

try:
    __version__ = str(pkg_resources.get_distribution("gpflow").parsed_version)
except pkg_resources.DistributionNotFound:
    __version__ = "develop"


# tests/examples/test_check_shapes_examples.py

# Copyright 2022 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.
"""
Code examples used in `check_shapes` docstrings.
"""


def test_example__basic() -> None:
    # This test / example is placed *above* the imports, to ensure that the example has the
    # necessary imports.

    # pylint: disable=import-outside-toplevel, redefined-outer-name, reimported

    # [basic]

    import tensorflow as tf

    from gpflow.experimental.check_shapes import check_shapes

    @tf.function
    @check_shapes(
        "features: [batch..., n_features]",
        "weights: [n_features]",
        "return: [batch...]",
    )
    def linear_model(features: tf.Tensor, weights: tf.Tensor) -> tf.Tensor:
        return tf.einsum("...i,i -> ...", features, weights)

    # [basic]

    w = tf.ones((3,))
    for batch_shape in [(), (2,), (2, 4)]:
        f = tf.ones(batch_shape + (3,))
        linear_model(f, w)


from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Iterable, Mapping, Optional, Sequence, Tuple
from unittest.mock import MagicMock

import numpy as np
import pytest

from gpflow.base import AnyNDArray
from gpflow.experimental.check_shapes import (
    DocstringFormat,
    ErrorContext,
    Shape,
    ShapeChecker,
    ShapeCheckingState,
    check_shape,
    check_shapes,
    disable_check_shapes,
    get_check_shapes,
    get_enable_check_shapes,
    get_enable_function_call_precompute,
    get_rewrite_docstrings,
    get_shape,
    inherit_check_shapes,
    register_get_shape,
    set_enable_check_shapes,
    set_enable_function_call_precompute,
    set_rewrite_docstrings,
)


def test_example__disable__manual() -> None:
    old_value = get_enable_check_shapes()

    try:

        # [disable__manual]

        set_enable_check_shapes(ShapeCheckingState.DISABLED)

        # [disable__manual]

    finally:
        set_enable_check_shapes(old_value)


def test_example__disable__context_manager() -> None:
    def performance_sensitive_function() -> None:
        pass

    # [disable__context_manager]

    with disable_check_shapes():
        performance_sensitive_function()

    # [disable__context_manager]


def test_example__pytest_fixture() -> None:
    # pylint: disable=unused-variable
    # [pytest_fixture]

    @pytest.fixture(autouse=True)
    def enable_shape_checks() -> Iterable[None]:
        old_enable = get_enable_check_shapes()
        old_rewrite_docstrings = get_rewrite_docstrings()
        old_function_call_precompute = get_enable_function_call_precompute()
        set_enable_check_shapes(ShapeCheckingState.ENABLED)
        set_rewrite_docstrings(DocstringFormat.SPHINX)
        set_enable_function_call_precompute(True)
        yield
        set_enable_function_call_precompute(old_function_call_precompute)
        set_rewrite_docstrings(old_rewrite_docstrings)
        set_enable_check_shapes(old_enable)

    # [pytest_fixture]


def test_example__argument_ref_attribute() -> None:
    # [argument_ref_attribute]

    @dataclass
    class Statistics:
        mean: AnyNDArray
        std: AnyNDArray

    @check_shapes(
        "data: [n_rows, n_columns]",
        "return.mean: [n_columns]",
        "return.std: [n_columns]",
    )
    def compute_statistics(data: AnyNDArray) -> Statistics:
        return Statistics(np.mean(data, axis=0), np.std(data, axis=0))

    # [argument_ref_attribute]

    compute_statistics(np.ones((4, 3)))


def test_example__argument_ref_index() -> None:
    # [argument_ref_index]

    @check_shapes(
        "data: [n_rows, n_columns]",
        "return[0]: [n_columns]",
        "return[1]: [n_columns]",
    )
    def compute_mean_and_std(data: AnyNDArray) -> Tuple[AnyNDArray, AnyNDArray]:
        return np.mean(data, axis=0), np.std(data, axis=0)

    # [argument_ref_index]

    compute_mean_and_std(np.ones((4, 3)))


def test_example__argument_ref_all() -> None:
    # [argument_ref_all]

    @check_shapes(
        "data[all]: [., n_columns]",
        "return: [., n_columns]",
    )
    def concat_rows(data: Sequence[AnyNDArray]) -> AnyNDArray:
        return np.concatenate(data, axis=0)

    concat_rows(
        [
            np.ones((1, 3)),
            np.ones((4, 3)),
        ]
    )

    # [argument_ref_all]


def test_example__argument_ref_keys() -> None:
    # [argument_ref_keys]

    @check_shapes(
        "data.keys(): [.]",
        "return: []",
    )
    def sum_key_lengths(data: Mapping[Tuple[int, ...], str]) -> int:
        return sum(len(k) for k in data)

    sum_key_lengths(
        {
            (3,): "foo",
            (1, 2): "bar",
        }
    )

    # [argument_ref_keys]


def test_example__argument_ref_values() -> None:
    # [argument_ref_values]

    @check_shapes(
        "data.values(): [., n_columns]",
        "return: [., n_columns]",
    )
    def concat_rows(data: Mapping[str, AnyNDArray]) -> AnyNDArray:
        return np.concatenate(list(data.values()), axis=0)

    concat_rows(
        {
            "foo": np.ones((1, 3)),
            "bar": np.ones((4, 3)),
        }
    )

    # [argument_ref_values]


def test_example__argument_ref_optional() -> None:
    # [argument_ref_optional]

    @check_shapes(
        "x1: [n_rows_1, n_inputs]",
        "x2: [n_rows_2, n_inputs]",
        "return: [n_rows_1, n_rows_2]",
    )
    def squared_exponential_kernel(
        variance: float, x1: AnyNDArray, x2: Optional[AnyNDArray] = None
    ) -> AnyNDArray:
        if x2 is None:
            x2 = x1
        cov: AnyNDArray = variance * np.exp(
            -0.5 * np.sum((x1[:, None, :] - x2[None, :, :]) ** 2, axis=2)
        )
        return cov

    squared_exponential_kernel(1.0, np.ones((3, 2)), np.ones((4, 2)))
    squared_exponential_kernel(3.2, np.ones((3, 2)))

    # [argument_ref_optional]


def test_example__dimension_spec_constant() -> None:
    # [dimension_spec_constant]

    @check_shapes(
        "v1: [2]",
        "v2: [2]",
    )
    def vector_2d_distance(v1: AnyNDArray, v2: AnyNDArray) -> float:
        return float(np.sqrt(np.sum((v1 - v2) ** 2)))

    # [dimension_spec_constant]

    vector_2d_distance(np.ones((2,)), np.ones((2,)))


def test_example__dimension_spec_variable() -> None:
    # [dimension_spec_variable]

    @check_shapes(
        "v1: [d]",
        "v2: [d]",
    )
    def vector_distance(v1: AnyNDArray, v2: AnyNDArray) -> float:
        return float(np.sqrt(np.sum((v1 - v2) ** 2)))

    # [dimension_spec_variable]

    vector_distance(np.ones((3,)), np.ones((3,)))


def test_example__dimension_spec_anonymous__dot() -> None:
    # [dimension_spec_anonymous__dot]

    @check_shapes(
        "v: [.]",
    )
    def vector_length(v: AnyNDArray) -> float:
        return float(np.sqrt(np.sum(v ** 2)))

    # [dimension_spec_anonymous__dot]

    vector_length(np.ones((3,)))


def test_example__dimension_spec_anonymous__none() -> None:
    # [dimension_spec_anonymous__none]

    @check_shapes(
        "v: [None]",
    )
    def vector_length(v: AnyNDArray) -> float:
        return float(np.sqrt(np.sum(v ** 2)))

    # [dimension_spec_anonymous__none]

    vector_length(np.ones((3,)))


def test_example__dimension_spec_variable_rank__star() -> None:
    # [dimension_spec_variable_rank__star]

    @check_shapes(
        "x: [*batch, n_columns]",
        "return: [*batch]",
    )
    def batch_mean(x: AnyNDArray) -> AnyNDArray:
        mean: AnyNDArray = np.mean(x, axis=-1)
        return mean

    # [dimension_spec_variable_rank__star]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape + (3,))
        batch_mean(x)


def test_example__dimension_spec_variable_rank__ellipsis() -> None:
    # [dimension_spec_variable_rank__ellipsis]

    @check_shapes(
        "x: [batch..., n_columns]",
        "return: [batch...]",
    )
    def batch_mean(x: AnyNDArray) -> AnyNDArray:
        mean: AnyNDArray = np.mean(x, axis=-1)
        return mean

    # [dimension_spec_variable_rank__ellipsis]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape + (3,))
        batch_mean(x)


def test_example__dimension_spec_anonymous_variable_rank__star() -> None:
    # [dimension_spec_anonymous_variable_rank__star]

    @check_shapes(
        "x: [*]",
    )
    def rank(x: AnyNDArray) -> int:
        return len(x.shape)

    # [dimension_spec_anonymous_variable_rank__star]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape)
        rank(x)


def test_example__dimension_spec_anonymous_variable_rank__ellipsis() -> None:
    # [dimension_spec_anonymous_variable_rank__ellipsis]

    @check_shapes(
        "x: [...]",
    )
    def rank(x: AnyNDArray) -> int:
        return len(x.shape)

    # [dimension_spec_anonymous_variable_rank__ellipsis]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape)
        rank(x)


def test_example__dimension_spec__scalar() -> None:
    # [dimension_spec__scalar]

    @check_shapes(
        "x: [...]",
        "return: []",
    )
    def mean(x: AnyNDArray) -> AnyNDArray:
        mean: AnyNDArray = np.sum(x) / x.size
        return mean

    # [dimension_spec__scalar]

    for batch_shape in [(), (2,), (2, 4)]:
        x = np.ones(batch_shape)
        mean(x)


def test_example__dimension_spec_broadcast() -> None:

    # [dimension_spec_broadcast]

    @check_shapes(
        "a: [broadcast batch...]",
        "b: [broadcast batch...]",
        "return: [batch...]",
    )
    def add(a: AnyNDArray, b: AnyNDArray) -> AnyNDArray:
        return a + b

    # [dimension_spec_broadcast]

    add(np.ones((3, 1)), np.ones((1, 4)))


def test_example__bool_spec_argument_ref() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_argument_ref]

    @check_shapes(
        "a: [broadcast batch...] if check_a",
        "b: [broadcast batch...] if check_b",
        "return: [batch...]",
    )
    def add(a: AnyNDArray, b: AnyNDArray, check_a: bool = True, check_b: bool = True) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)

    # [bool_spec_argument_ref]


def test_example__bool_spec_argument_ref_is_none() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_argument_ref_is_none]

    @check_shapes(
        "a: [n_a]",
        "b: [n_b]",
        "return: [n_a, n_a] if b is None",
        "return: [n_a, n_b] if b is not None",
    )
    def square(a: AnyNDArray, b: Optional[AnyNDArray] = None) -> AnyNDArray:
        if b is None:
            b = a
        result: AnyNDArray = a[:, None] * b[None, :]
        return result

    square(np.ones((3,)))
    square(np.ones((3,)), np.ones((4,)))

    # [bool_spec_argument_ref_is_none]


def test_example__bool_spec_or() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_or]

    @check_shapes(
        "a: [broadcast batch...] if check_all or check_a",
        "b: [broadcast batch...] if check_all or check_b",
        "return: [batch...]",
    )
    def add(
        a: AnyNDArray,
        b: AnyNDArray,
        check_all: bool = False,
        check_a: bool = True,
        check_b: bool = True,
    ) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)

    # [bool_spec_or]


def test_example__bool_spec_and() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_and]

    @check_shapes(
        "a: [broadcast batch...] if enable_checks and check_a",
        "b: [broadcast batch...] if enable_checks and check_b",
        "return: [batch...]",
    )
    def add(
        a: AnyNDArray,
        b: AnyNDArray,
        enable_checks: bool = True,
        check_a: bool = True,
        check_b: bool = True,
    ) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)), check_b=False)

    # [bool_spec_and]


def test_example__bool_spec_not() -> None:
    # pylint: disable=unused-argument

    # [bool_spec_not]

    @check_shapes(
        "a: [broadcast batch...] if not disable_checks",
        "b: [broadcast batch...] if not disable_checks",
        "return: [batch...]",
    )
    def add(a: AnyNDArray, b: AnyNDArray, disable_checks: bool = False) -> AnyNDArray:
        return a + b

    add(np.ones((3, 1)), np.ones((1, 4)))

    # [bool_spec_not]


def test_example__bool_spec__composition() -> None:
    # pylint: disable=unused-argument

    # [bool_spec__composition]

    @check_shapes(
        "a: [j] if a_vector",
        "a: [i, j] if (not a_vector)",
        "b: [j] if b_vector",
        "b: [j, k] if (not b_vector)",
        "return: [1, 1] if a_vector and b_vector",
        "return: [1, k] if a_vector and (not b_vector)",
        "return: [i, 1] if (not a_vector) and b_vector",
        "return: [i, k] if (not a_vector) and (not b_vector)",
    )
    def multiply(a: AnyNDArray, b: AnyNDArray, a_vector: bool, b_vector: bool) -> AnyNDArray:
        if a_vector:
            a = a[None, :]
        if b_vector:
            b = b[:, None]

        return a @ b

    multiply(np.ones((4,)), np.ones((4, 5)), a_vector=True, b_vector=False)

    # [bool_spec__composition]

    multiply(np.ones((4,)), np.ones((4,)), a_vector=True, b_vector=True)
    multiply(np.ones((3, 4)), np.ones((4,)), a_vector=False, b_vector=True)
    multiply(np.ones((3, 4)), np.ones((4, 5)), a_vector=False, b_vector=False)


def test_example__note_spec__global() -> None:
    # [note_spec__global]

    @check_shapes(
        "features: [batch..., n_features]",
        "# linear_model currently only supports a single output.",
        "weights: [n_features]",
        "return: [batch...]",
    )
    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        prediction: AnyNDArray = np.einsum("...i,i -> ...", features, weights)
        return prediction

    # [note_spec__global]

    linear_model(np.ones((4, 3)), np.ones((3,)))


def test_example__note_spec__local() -> None:
    # [note_spec__local]

    @check_shapes(
        "features: [batch..., n_features]",
        "weights: [n_features] # linear_model currently only supports a single output.",
        "return: [batch...]",
    )
    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        prediction: AnyNDArray = np.einsum("...i,i -> ...", features, weights)
        return prediction

    # [note_spec__local]

    linear_model(np.ones((4, 3)), np.ones((3,)))


def test_example__reuse__inherit_check_shapes() -> None:
    # [reuse__inherit_check_shapes]

    class Model(ABC):
        @abstractmethod
        @check_shapes(
            "features: [batch..., n_features]",
            "return: [batch...]",
        )
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            pass

    class LinearModel(Model):
        @check_shapes(
            "weights: [n_features]",
        )
        def __init__(self, weights: AnyNDArray) -> None:
            self._weights = weights

        @inherit_check_shapes
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            prediction: AnyNDArray = np.einsum("...i,i -> ...", features, self._weights)
            return prediction

    # [reuse__inherit_check_shapes]

    model = LinearModel(np.ones((3,)))
    model.predict(np.ones((10, 3)))


def test_example__reuse__functional() -> None:
    # [reuse__functional]

    check_metric_shapes = check_shapes(
        "actual: [n_rows, n_labels]",
        "predicted: [n_rows, n_labels]",
        "return: []",
    )

    @check_metric_shapes
    def rmse(actual: AnyNDArray, predicted: AnyNDArray) -> float:
        return float(np.mean(np.sqrt(np.mean((predicted - actual) ** 2, axis=-1))))

    @check_metric_shapes
    def mape(actual: AnyNDArray, predicted: AnyNDArray) -> float:
        return float(np.mean(np.abs((predicted - actual) / actual)))

    # [reuse__functional]

    actual = np.ones((10, 3))
    predicted = np.ones((10, 3))
    rmse(actual, predicted)
    mape(actual, predicted)


def test_example__reuse__get_check_shapes() -> None:
    # [reuse__get_check_shapes]

    class Model(ABC):
        @abstractmethod
        @check_shapes(
            "features: [batch..., n_features]",
            "return: [batch...]",
        )
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            pass

    @check_shapes(
        "test_features: [n_rows, n_features]",
        "test_labels: [n_rows]",
    )
    def evaluate_model(model: Model, test_features: AnyNDArray, test_labels: AnyNDArray) -> float:
        prediction = model.predict(test_features)
        return float(np.mean(np.sqrt(np.mean((prediction - test_labels) ** 2, axis=-1))))

    def test_evaluate_model() -> None:
        fake_features = np.ones((10, 3))
        fake_labels = np.ones((10,))
        fake_predictions = np.ones((10,))

        @get_check_shapes(Model.predict)
        def fake_predict(features: AnyNDArray) -> AnyNDArray:
            assert features is fake_features
            return fake_predictions

        fake_model = MagicMock(spec=Model, predict=fake_predict)

        assert pytest.approx(0.0) == evaluate_model(fake_model, fake_features, fake_labels)

    # [reuse__get_check_shapes]

    test_evaluate_model()


def test_example__intermediate_results() -> None:
    # [intermediate_results]

    @check_shapes(
        "weights: [n_features, n_labels]",
        "test_features: [n_rows, n_features]",
        "test_labels: [n_rows, n_labels]",
        "return: []",
    )
    def loss(weights: AnyNDArray, test_features: AnyNDArray, test_labels: AnyNDArray) -> AnyNDArray:
        prediction = check_shape(test_features @ weights, "[n_rows, n_labels]")
        error: AnyNDArray = check_shape(prediction - test_labels, "[n_rows, n_labels]")
        square_error = check_shape(error ** 2, "[n_rows, n_labels]")
        mean_square_error = check_shape(np.mean(square_error, axis=-1), "[n_rows]")
        root_mean_square_error = check_shape(np.sqrt(mean_square_error), "[n_rows]")
        loss: AnyNDArray = np.mean(root_mean_square_error)
        return loss

    # [intermediate_results]

    loss(
        np.ones((5, 3)),
        np.ones((10, 5)),
        np.ones((10, 3)),
    )


def test_example__shape_checker__raw() -> None:
    # [shape_checker__raw]

    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        checker = ShapeChecker()
        checker.check_shape(features, "[batch..., n_features]")
        checker.check_shape(weights, "[n_features]")
        prediction: AnyNDArray = checker.check_shape(
            np.einsum("...i,i -> ...", features, weights), "[batch...]"
        )
        return prediction

    # [shape_checker__raw]

    linear_model(np.ones((4, 3)), np.ones((3,)))


def test_example__disable_function_call_precompute() -> None:
    def buggy_function() -> None:
        pass

    old_value = get_enable_function_call_precompute()

    try:

        # [disable_function_call_precompute]

        set_enable_function_call_precompute(True)

        buggy_function()

        # [disable_function_call_precompute]

    finally:
        set_enable_function_call_precompute(old_value)


def test_example__doc_rewrite() -> None:

    # [doc_rewrite__definition]

    @check_shapes(
        "features: [batch..., n_features]",
        "weights: [n_features]",
        "return: [batch...]",
    )
    def linear_model(features: AnyNDArray, weights: AnyNDArray) -> AnyNDArray:
        """
        Computes a prediction from a linear model.

        :param features: Data to make predictions from.
        :param weights: Model weights.
        :returns: Model predictions.
        """
        prediction: AnyNDArray = np.einsum("...i,i -> ...", features, weights)
        return prediction

    # [doc_rewrite__definition]

    expected_doc = (
        # [doc_rewrite__rewritten]
        """
        Computes a prediction from a linear model.

        :param features:
            * **features** has shape [*batch*..., *n_features*].

            Data to make predictions from.
        :param weights:
            * **weights** has shape [*n_features*].

            Model weights.
        :returns:
            * **return** has shape [*batch*...].

            Model predictions.
        """
        # [doc_rewrite__rewritten]
    )

    assert expected_doc == linear_model.__doc__


def test_example__doc_rewrite__disable() -> None:

    old_rewrite_docstrings = get_rewrite_docstrings()

    try:
        # [doc_rewrite__disable]
        set_rewrite_docstrings(None)
        # [doc_rewrite__disable]
    finally:
        set_rewrite_docstrings(old_rewrite_docstrings)


def test_example__custom_type() -> None:
    # pylint: disable=protected-access,unused_variable

    # [custom_type]

    class LinearModel:
        @check_shapes(
            "weights: [n_features]",
        )
        def __init__(self, weights: AnyNDArray) -> None:
            self._weights = weights

        @check_shapes(
            "self: [n_features]",
            "features: [batch..., n_features]",
            "return: [batch...]",
        )
        def predict(self, features: AnyNDArray) -> AnyNDArray:
            prediction: AnyNDArray = np.einsum("...i,i -> ...", features, self._weights)
            return prediction

    @register_get_shape(LinearModel)
    def get_linear_model_shape(model: LinearModel, context: ErrorContext) -> Shape:
        shape: Shape = model._weights.shape
        return shape

    @check_shapes(
        "model: [n_features]",
        "test_features: [n_rows, n_features]",
        "test_labels: [n_rows]",
        "return: []",
    )
    def loss(model: LinearModel, test_features: AnyNDArray, test_labels: AnyNDArray) -> float:
        prediction = model.predict(test_features)
        return float(np.mean(np.sqrt(np.mean((prediction - test_labels) ** 2, axis=-1))))

    # [custom_type]

    model = LinearModel(np.ones((3,)))
    loss(model, np.ones((10, 3)), np.ones((10,)))


# tests/gpflow/conftest.py

#  Copyright 2021 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 typing import cast

import numpy as np
import pytest
from _pytest.fixtures import SubRequest


@pytest.fixture(name="full_cov", params=[True, False])
def _full_cov_fixture(request: SubRequest) -> bool:
    return cast(bool, request.param)


@pytest.fixture(name="full_output_cov", params=[True, False])
def _full_output_cov_fixture(request: SubRequest) -> bool:
    return cast(bool, request.param)


@pytest.fixture(name="whiten", params=[True, False])
def _whiten_fixture(request: SubRequest) -> bool:
    return cast(bool, request.param)


@pytest.fixture(name="rng")
def _rng_fixture() -> np.random.Generator:
    return np.random.default_rng(20220523)


# tests/gpflow/test_all.py

# Copyright 2022 the GPflow authors.
#
# 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.
import ast
import importlib
from pathlib import Path
from types import ModuleType
from typing import Any, List, Optional, Sequence

import pytest

import gpflow


def is_dunder(name: str) -> bool:
    return name.startswith("__") and name.endswith("__")


def get_module_path(module: ModuleType) -> Path:
    path_str = module.__file__
    assert path_str is not None
    return Path(path_str)


def is_package(module: ModuleType) -> bool:
    return get_module_path(module).name == "__init__.py"


def find_modules(root: ModuleType) -> Sequence[ModuleType]:
    root_path = get_module_path(root).parent
    parent = root_path.parent
    result: List[ModuleType] = []
    for path in root_path.glob("**/*.py"):
        relative_path = path.relative_to(parent)
        if relative_path.name == "__init__.py":
            module_name__slashes = str(relative_path.parent)
        else:
            module_name__slashes = str(relative_path)[:-3]
        module_name = module_name__slashes.replace("/", ".")
        result.append(importlib.import_module(module_name))
    return result


_MODULES = find_modules(gpflow)
_PACKAGES = [m for m in _MODULES if is_package(m)]
_MODULES_WITH_ALL = [m for m in _MODULES if hasattr(m, "__all__")]


@pytest.mark.parametrize("package", _PACKAGES)
def test_all_present_and_up_to_date(package: ModuleType) -> None:
    imported_sorted = sorted(attr for attr in dir(package) if not is_dunder(attr))
    all_list = getattr(package, "__all__", None)
    assert all_list is not None, f"Package {package} is missing an explicit __all__."
    all_sorted = sorted(attr for attr in all_list if not is_dunder(attr))
    assert imported_sorted == all_sorted, (
        f"{package}.__all__ is outdated."
        f" Imported values are {imported_sorted}, but exported values are {all_list}."
    )


@pytest.mark.parametrize("module", _MODULES_WITH_ALL)
def test_all_sorted(module: ModuleType) -> None:
    all_list = getattr(module, "__all__", None)
    assert all_list is not None  # Hint for mypy.
    all_sorted = sorted(all_list)
    assert (
        all_sorted == all_list
    ), f"{module}.__all__ is not sorted. Expected {all_sorted}, found {all_list}."


@pytest.mark.parametrize("module", _MODULES_WITH_ALL)
def test_all_static(module: ModuleType) -> None:
    module_path = get_module_path(module)
    tree = ast.parse(module_path.read_text(), str(module_path))
    assert isinstance(tree, ast.Module)

    def is_all_assignment(t: Any) -> bool:
        return (
            isinstance(t, ast.Assign)
            and len(t.targets) == 1
            and isinstance(t.targets[0], ast.Name)
            and t.targets[0].id == "__all__"
        )

    all_assignment: Optional[ast.Assign] = None
    for t in tree.body:
        if is_all_assignment(t):
            all_assignment = t  # type: ignore[assignment]
    assert all_assignment is not None, f"No `__all__` found in module {module}."

    error_message = "__all__ must be a static list of constant strings. Some tools expect this."

    all_value = all_assignment.value
    assert isinstance(all_value, ast.List), error_message

    assert all(
        # Actual type depends on Python version:
        (isinstance(t, ast.Constant) or isinstance(t, ast.Str))
        for t in all_value.elts
    ), error_message


# tests/gpflow/test_base.py

# Copyright 2020 the GPflow authors.
#
# 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.


import tempfile
from typing import Sequence

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp

import gpflow
from gpflow.base import AnyNDArray, PriorOn, TensorData
from gpflow.config import default_float
from gpflow.utilities import positive, triangular


def test_parameter_assign_validation() -> None:
    with pytest.raises(tf.errors.InvalidArgumentError):
        param = gpflow.Parameter(0.0, transform=positive())

    param = gpflow.Parameter(0.1, transform=positive())
    param.assign(0.2)
    with pytest.raises(tf.errors.InvalidArgumentError):
        param.assign(0.0)


def test_cast_to_dtype_precision_issue() -> None:
    """
    TensorFlow's tf.cast(value, dtype) implicitly does a tf.convert_to_tensor(value)
    *before* the cast when the value is not a tensor already. When value is a python float,
    this results in the following behaviour:

    >>> tf.cast(0.2, tf.float64)
    <tf.Tensor: id=37, shape=(), dtype=float64, numpy=0.20000000298023224>

    instead of the expected expansion of 0.2 to float64 precision that you get when
    passing in an object that already carries dtype information, such as a numpy array
    (which has float64 precision by default):

    >>> tf.cast(np.array(0.2), tf.float64)
    <tf.Tensor: id=40, shape=(), dtype=float64, numpy=0.2>

    This affected *all* gpflow.Parameter objects, resulting in numerical discrepancies
    between GPflow 1 and 2, due to the pass through _cast_to_dtype, which is now fixed.
    This is the corresponding regression test.
    """
    p = gpflow.Parameter(0.2, dtype=np.float64)
    actual_value = p.numpy()
    assert actual_value.dtype == np.float64
    expected_value = np.float64(0.2)
    assert actual_value == expected_value


def test_parameter_saved() -> None:
    dtype = tf.float64

    class Model(tf.Module):
        def __init__(self) -> None:
            self.p = gpflow.Parameter(0.1, dtype=dtype, transform=gpflow.utilities.positive())

        @tf.function(input_signature=[tf.TensorSpec([], dtype=dtype)])
        def exec(self, x: tf.Tensor) -> tf.Tensor:
            return tf.square(x * self.p)

    m0 = Model()
    x = tf.convert_to_tensor(2.0, dtype=dtype)
    expected = m0.exec(x)
    with tempfile.TemporaryDirectory() as dirname:
        tf.saved_model.save(m0, dirname)
        m1 = tf.saved_model.load(dirname)
        actual = m1.exec(x)
        np.testing.assert_equal(actual, expected)


@pytest.mark.parametrize("value", [0.0, [1.2, 1.1]])
def test_construct_parameter_from_existing_parameter_check_value(value: TensorData) -> None:
    initial_parameter = gpflow.Parameter(value)
    new_parameter = gpflow.Parameter(initial_parameter)

    np.testing.assert_equal(new_parameter.numpy(), value)


@pytest.mark.parametrize("value", [0.0, [1.2, 1.1]])
def test_construct_parameter_from_existing_parameter_override_value(value: TensorData) -> None:
    initial_parameter = gpflow.Parameter(value)
    new_parameter = gpflow.Parameter(initial_parameter + 1.0)

    np.testing.assert_equal(new_parameter.numpy(), np.array(value) + 1.0)


def test_construct_parameter_from_existing_parameter_check_transform() -> None:
    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    initial_parameter = gpflow.Parameter([1.2, 1.1], transform=transform)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.transform == transform


def test_construct_parameter_from_existing_parameter_override_transform() -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])

    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    new_parameter = gpflow.Parameter(initial_parameter, transform=transform)

    assert new_parameter.transform == transform


def test_construct_parameter_from_existing_parameter_check_prior() -> None:
    prior = tfp.distributions.Normal(0.0, 1.0)
    initial_parameter = gpflow.Parameter([1.2, 1.1], prior=prior)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.prior == prior


def test_construct_parameter_from_existing_parameter_override_prior() -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])

    prior = tfp.distributions.Normal(0.0, 1.0)
    new_parameter = gpflow.Parameter(initial_parameter, prior=prior)

    assert new_parameter.prior == prior


@pytest.mark.parametrize("prior_on", [PriorOn.CONSTRAINED, PriorOn.UNCONSTRAINED])
def test_construct_parameter_from_existing_parameter_check_prior_on(prior_on: PriorOn) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1], prior_on=prior_on)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.prior_on == prior_on


@pytest.mark.parametrize("prior_on", [PriorOn.CONSTRAINED, PriorOn.UNCONSTRAINED])
def test_construct_parameter_from_existing_parameter_override_prior_on(prior_on: PriorOn) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])
    new_parameter = gpflow.Parameter(initial_parameter, prior_on=prior_on)

    assert new_parameter.prior_on == prior_on


@pytest.mark.parametrize("trainable", [True, False])
def test_construct_parameter_from_existing_parameter_check_trainable(trainable: bool) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1], trainable=trainable)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.trainable == trainable


@pytest.mark.parametrize("trainable", [True, False])
def test_construct_parameter_from_existing_parameter_override_trainable(trainable: bool) -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1], trainable=trainable)
    new_parameter = gpflow.Parameter(initial_parameter, trainable=not trainable)

    assert new_parameter.trainable is not trainable


@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
def test_construct_parameter_from_existing_parameter_check_dtype(dtype: tf.DType) -> None:
    initial_parameter = gpflow.Parameter([1.1, 2.1], dtype=dtype)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.dtype == dtype


@pytest.mark.parametrize("dtype", [tf.float32, tf.float64])
def test_construct_parameter_from_existing_parameter_override_dtype(dtype: tf.DType) -> None:
    initial_parameter = gpflow.Parameter([1.1, 2.1])
    new_parameter = gpflow.Parameter(initial_parameter, dtype=dtype)

    assert new_parameter.dtype == dtype


def test_construct_parameter_from_existing_parameter_check_name() -> None:
    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    initial_parameter = gpflow.Parameter([1.2, 1.1], transform=transform)
    new_parameter = gpflow.Parameter(initial_parameter)

    assert new_parameter.name == transform.name


def test_construct_parameter_from_existing_parameter_override_name() -> None:
    initial_parameter = gpflow.Parameter([1.2, 1.1])
    transform = tfp.bijectors.Sigmoid(
        tf.constant(0.0, dtype=tf.float64), tf.constant(2.0, dtype=tf.float64)
    )
    new_parameter = gpflow.Parameter(initial_parameter, transform=transform)

    assert new_parameter.name == transform.name


def test_construct_parameter_from_existing_parameter_value_becomes_invalid() -> None:
    initial_parameter = gpflow.Parameter(0.0)
    transform = tfp.bijectors.Reciprocal()

    with pytest.raises(tf.errors.InvalidArgumentError) as exc:
        gpflow.Parameter(initial_parameter, transform=transform)

    assert "gpflow.Parameter" in exc.value.message


def test_construct_parameter_with_variable_shape() -> None:
    parameter = gpflow.Parameter([[1, 2, 3]], shape=[None, None])

    values: Sequence[AnyNDArray] = [
        np.ones((0, 0), dtype=default_float()),
        np.ones((1, 3), dtype=default_float()),
        np.ones((3, 1), dtype=default_float()),
        np.ones((3, 4), dtype=default_float()),
    ]

    for value in values:
        parameter.assign(value)
        np.testing.assert_equal(value, parameter.numpy())


def test_construct_parameter_with_variable_shape__different_constrained_shape() -> None:
    parameter = gpflow.Parameter(
        [[1, 0], [2, 3]],
        transform=triangular(),
        unconstrained_shape=[None],
        constrained_shape=[None, None],
    )

    values: Sequence[AnyNDArray] = [
        # The triangular() transform doesn't appear to support 0x0 matrices.
        np.tril(np.ones((1, 1), dtype=default_float())),
        np.tril(np.ones((2, 2), dtype=default_float())),
        np.tril(np.ones((3, 3), dtype=default_float())),
    ]

    for value in values:
        parameter.assign(value)
        np.testing.assert_equal(value, parameter.numpy())


# tests/gpflow/test_base_prior.py

from typing import Any, Tuple, Type

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.bijectors import Exp
from tensorflow_probability.python.distributions import Uniform

import gpflow
from gpflow.base import AnyNDArray, PriorOn
from gpflow.config import set_default_float
from gpflow.utilities import to_default_float

np.random.seed(1)


class Datum:
    X: AnyNDArray = 10 * np.random.randn(5, 1)
    Y: AnyNDArray = 10 * np.random.randn(5, 1)
    lengthscale = 3.3


def test_gpr_objective_equivalence() -> None:
    """
    In Maximum Likelihood Estimation (MLE), i.e. when there are no priors on
    the parameters, the objective should not depend on any transforms on the
    parameters.
    We use GPR as a simple model that has an objective.
    """
    data = (Datum.X, Datum.Y)
    l_value = Datum.lengthscale

    l_variable = tf.Variable(l_value, dtype=gpflow.default_float(), trainable=True)
    m1 = gpflow.models.GPR(data, kernel=gpflow.kernels.SquaredExponential(lengthscales=l_value))
    m2 = gpflow.models.GPR(data, kernel=gpflow.kernels.SquaredExponential())
    m2.kernel.lengthscales = gpflow.Parameter(l_variable, transform=None)
    assert np.allclose(
        m1.kernel.lengthscales.numpy(), m2.kernel.lengthscales.numpy()
    )  # consistency check

    assert np.allclose(
        m1.log_marginal_likelihood().numpy(), m2.log_marginal_likelihood().numpy()
    ), "MLE objective should not depend on Parameter transform"


def test_log_prior_with_no_prior() -> None:
    """
    A parameter without any prior should have zero log-prior,
    even if it has a transform to constrain it.
    """
    param = gpflow.Parameter(5.3, transform=gpflow.utilities.positive())
    assert param.log_prior_density().numpy() == 0.0


def test_log_prior_for_uniform_prior() -> None:
    """
    If we assign a Uniform prior to a parameter, we should not expect the value of the prior density
    to change with the parameter value, even if it has a transform associated with it.
    """

    uniform_prior = Uniform(low=np.float64(0), high=np.float64(100))
    param = gpflow.Parameter(1.0, transform=gpflow.utilities.positive(), prior=uniform_prior)
    low_value = param.log_prior_density().numpy()
    param.assign(10.0)
    high_value = param.log_prior_density().numpy()

    assert np.isclose(low_value, high_value)


def test_log_prior_on_unconstrained() -> None:
    """
    A parameter with an Exp transform, and a uniform prior on its unconstrained, should have a
    prior in the constrained space that scales as 1/value.
    """

    initial_value = 1.0
    scale_factor = 10.0
    uniform_prior = Uniform(low=np.float64(0), high=np.float64(100))
    param = gpflow.Parameter(
        initial_value,
        transform=Exp(),
        prior=uniform_prior,
        prior_on=PriorOn.UNCONSTRAINED,
    )
    low_value = param.log_prior_density().numpy()
    param.assign(scale_factor * initial_value)
    high_value = param.log_prior_density().numpy()

    assert np.isclose(low_value, high_value + np.log(scale_factor))


class DummyModel(gpflow.models.BayesianModel):
    value = 3.3
    log_scale = 0.4

    def __init__(self, with_transform: bool) -> None:
        super().__init__()

        prior = tfp.distributions.Normal(to_default_float(1.0), to_default_float(1.0))

        scale = np.exp(self.log_scale)
        if with_transform:
            transform = tfp.bijectors.Shift(to_default_float(0.0))(
                tfp.bijectors.Scale(to_default_float(scale))
            )
        else:
            transform = None

        self.theta = gpflow.Parameter(self.value, prior=prior, transform=transform)

    def maximum_log_likelihood_objective(self, *args: Any, **kwargs: Any) -> tf.Tensor:
        assert not args
        assert not kwargs
        return (self.theta + 5) ** 2


def test_map_invariance_to_transform() -> None:
    m1 = DummyModel(with_transform=True)
    m2 = DummyModel(with_transform=False)
    assert np.allclose(
        m1.log_posterior_density().numpy(), m2.log_posterior_density().numpy()
    ), "log posterior density should not be affected by a transform"


def get_gpmc_model_params() -> Tuple[Any, ...]:
    kernel = gpflow.kernels.Matern32()
    likelihood = gpflow.likelihoods.Gaussian()
    data = [np.random.randn(5, 1), np.random.randn(5, 1)]
    return data, kernel, likelihood


@pytest.mark.parametrize(
    "model_class, args",
    [
        (gpflow.models.GPMC, get_gpmc_model_params()),
        # (gpflow.models.SGPMC, get_SGPMC_model_params()) # Fails due to inducing_variable=None bug
    ],
)
def test_v_prior_dtypes(model_class: Type[Any], args: Tuple[Any, ...]) -> None:
    with gpflow.config.as_context():
        set_default_float(np.float32)
        m = model_class(*args)
        assert m.V.prior.dtype == np.float32
        set_default_float(np.float64)
        m = model_class(*args)
        assert m.V.prior.dtype == np.float64


# tests/gpflow/test_base_training.py

# Copyright 2018 the GPflow authors.
#
# 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 typing import Any, Mapping

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.models import SVGP
from gpflow.utilities import leaf_components, multiple_assign, read_values

rng = np.random.RandomState(0)


class Data:
    H0 = 5
    H1 = 2
    M = 10
    D = 1
    Z = rng.rand(M, 1)
    ls = 2.0
    ls_new = 1.5
    var = 1.0
    var_new = 10.0


# ------------------------------------------
# Fixtures
# ------------------------------------------


@pytest.fixture
def model() -> SVGP:
    return SVGP(
        kernel=gpflow.kernels.SquaredExponential(lengthscales=Data.ls, variance=Data.var),
        likelihood=gpflow.likelihoods.Gaussian(),
        inducing_variable=Data.Z,
        q_diag=True,
    )


# ------------------------------------------
# Reference
# ------------------------------------------

model_param_updates = {
    ".kernel.lengthscales": Data.ls_new,
    ".likelihood.variance": Data.var_new,
    ".inducing_variable.Z": np.zeros_like(Data.Z),
    ".q_sqrt": 0.5 * np.ones((Data.M, 1)),
}
model_wrong_path = [
    {"kernel.lengthscales": Data.ls_new},
    {".Gaussian.variance": Data.var_new},
    {"inducing_variable.Z": np.zeros_like(Data.Z)},
    {".q_std": 0.5 * np.ones((Data.M, 1))},
]

model_wrong_value = [
    {".likelihood.variance": np.ones((2, 1), dtype=np.int32)},
    {".inducing_variable.Z": [1, 2, 3]},
]


@pytest.mark.parametrize("var_update_dict", [model_param_updates])
def test_multiple_assign_updates_correct_values(
    model: SVGP, var_update_dict: Mapping[str, Any]
) -> None:
    old_value_dict = leaf_components(model)
    multiple_assign(model, var_update_dict)
    for path, variable in leaf_components(model).items():
        if path in var_update_dict.keys():
            np.testing.assert_almost_equal(variable.numpy(), var_update_dict[path], decimal=7)
        else:
            np.testing.assert_equal(variable.numpy(), old_value_dict[path].numpy())


@pytest.mark.parametrize("wrong_var_update_dict", model_wrong_path)
def test_multiple_assign_fails_with_invalid_path(
    model: SVGP, wrong_var_update_dict: Mapping[str, Any]
) -> None:
    with pytest.raises(KeyError):
        multiple_assign(model, wrong_var_update_dict)


@pytest.mark.parametrize("wrong_var_update_dict", model_wrong_value)
def test_multiple_assign_fails_with_invalid_values(
    model: SVGP, wrong_var_update_dict: Mapping[str, Any]
) -> None:
    with pytest.raises(ValueError):
        multiple_assign(model, wrong_var_update_dict)


def test_dict_utilities(model: SVGP) -> None:
    """
    Test both `parameter_dict()` and `read_values()`
    """

    class SubModule(tf.Module):
        def __init__(self) -> None:
            self.parameter = gpflow.Parameter(1.0)
            self.variable = tf.Variable(1.0)

    class Module(tf.Module):
        def __init__(self) -> None:
            self.submodule = SubModule()
            self.top_parameter = gpflow.Parameter(3.0)

    m = Module()
    params = gpflow.utilities.parameter_dict(m)
    # {
    #   ".submodule.parameter": <parameter object>,
    #   ".submodule.variable": <variable object>
    # }
    assert list(params.keys()) == [
        ".submodule.parameter",
        ".submodule.variable",
        ".top_parameter",
    ]
    assert list(params.values()) == [
        m.submodule.parameter,
        m.submodule.variable,
        m.top_parameter,
    ]

    for k, v in read_values(m).items():
        assert params[k].numpy() == v


# tests/gpflow/test_functions.py

# Copyright 2017 the GPflow authors.
#
# 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 typing import Any, Sequence, Tuple, Type

import numpy as np
import pytest
from numpy.testing import assert_allclose

import gpflow
from gpflow.base import AnyNDArray, TensorType
from gpflow.config import default_int
from gpflow.experimental.check_shapes import check_shapes
from gpflow.functions import (
    Additive,
    Constant,
    Linear,
    MeanFunction,
    Polynomial,
    Product,
    SwitchedMeanFunction,
    Zero,
)
from gpflow.inducing_variables import InducingPoints

rng = np.random.RandomState(99021)


class Datum:
    input_dim, output_dim = 3, 2
    N, Ntest, M = 20, 30, 10


_mean_functions = [
    Zero(),
    Linear(
        A=rng.randn(Datum.input_dim, Datum.output_dim),
        b=rng.randn(Datum.output_dim, 1).reshape(-1),
    ),
    Constant(c=rng.randn(Datum.output_dim, 1).reshape(-1)),
    Polynomial(degree=2, input_dim=Datum.input_dim, output_dim=Datum.output_dim),
]


@pytest.mark.parametrize("mean_function_1", _mean_functions)
@pytest.mark.parametrize("mean_function_2", _mean_functions)
@pytest.mark.parametrize("operation", ["+", "*"])
def test_mean_functions_output_shape(
    mean_function_1: MeanFunction, mean_function_2: MeanFunction, operation: str
) -> None:
    """
    Test the output shape for basic and compositional mean functions, also
    check that the combination of mean functions returns the correct class
    """
    X = np.random.randn(Datum.N, Datum.input_dim)
    Y = mean_function_1(X)
    # basic output shape check
    assert Y.shape in [(Datum.N, Datum.output_dim), (Datum.N, 1)]

    # composed mean function output shape check
    if operation == "+":
        mean_composed = mean_function_1 + mean_function_2
    elif operation == "*":
        mean_composed = mean_function_1 * mean_function_2
    else:
        raise NotImplementedError

    Y_composed = mean_composed(X)
    assert Y_composed.shape in [(Datum.N, Datum.output_dim), (Datum.N, 1)]


@pytest.mark.parametrize("mean_function_1", _mean_functions)
@pytest.mark.parametrize("mean_function_2", _mean_functions)
@pytest.mark.parametrize("operation", ["+", "*"])
def test_mean_functions_composite_type(
    mean_function_1: MeanFunction, mean_function_2: MeanFunction, operation: str
) -> None:
    if operation == "+":
        mean_composed = mean_function_1 + mean_function_2
        assert isinstance(mean_composed, Additive)
    elif operation == "*":
        mean_composed = mean_function_1 * mean_function_2
        assert isinstance(mean_composed, Product)
    else:
        raise NotImplementedError


_linear_functions = [
    Linear(
        A=rng.randn(Datum.input_dim, Datum.output_dim),
        b=rng.randn(Datum.output_dim, 1).reshape(-1),
    )
    for _ in range(3)
]

# Append inverse of first Linear mean function in _linear_functions
_linear_functions.append(Linear(A=-1.0 * _linear_functions[0].A, b=-1.0 * _linear_functions[0].b))

_constant_functions = [Constant(c=rng.randn(Datum.output_dim, 1).reshape(-1)) for _ in range(3)]
# Append inverse of first Constant mean function in _constant_functions
_constant_functions.append(Constant(c=-1.0 * _constant_functions[0].c))


@check_shapes(
    "X: [N, D]",
    "Y: [N, Q]",
)
def _create_GPR_model_with_bias(
    X: TensorType, Y: TensorType, mean_function: MeanFunction
) -> gpflow.models.GPR:
    return gpflow.models.GPR(
        (X, Y), mean_function=mean_function, kernel=gpflow.kernels.Bias(Datum.input_dim)
    )


@pytest.mark.parametrize("mean_functions", [_linear_functions, _constant_functions])
def test_mean_functions_distributive_property(mean_functions: Sequence[MeanFunction]) -> None:
    """
    Tests that distributive property of addition and multiplication holds for mean functions
    (both Constant and Linear): A * (B + C) = A * B + A * C
    """
    X, Y = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, Datum.output_dim)
    Xtest = rng.randn(30, Datum.input_dim)
    A, B, C = mean_functions[0], mean_functions[1], mean_functions[2]
    lhs = Product(A, Additive(B, C))  # A * (B + C)
    rhs = Additive(Product(A, B), Product(A, C))  # A * B + A * C

    model_lhs = _create_GPR_model_with_bias(X, Y, mean_function=lhs)
    model_rhs = _create_GPR_model_with_bias(X, Y, mean_function=rhs)

    mu_lhs, var_lhs = model_lhs.predict_f(Xtest)
    mu_rhs, var_rhs = model_rhs.predict_f(Xtest)

    assert_allclose(mu_lhs, mu_rhs)
    assert_allclose(var_lhs, var_rhs)


@pytest.mark.parametrize("mean_functions", [_linear_functions, _constant_functions])
def test_mean_functions_A_minus_A_equals_zero(mean_functions: Sequence[MeanFunction]) -> None:
    """
    Tests that the addition the inverse of a mean function to itself is equivalent to having a
    Zero mean function: A + (-A) = 0
    """
    X, Y = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, Datum.output_dim)
    Xtest = rng.randn(30, Datum.input_dim)
    A, A_inverse = mean_functions[0], mean_functions[-1]
    lhs = Additive(A, A_inverse)  # A + (-A)
    rhs = Zero()  # 0

    model_lhs = _create_GPR_model_with_bias(X, Y, mean_function=lhs)
    model_rhs = _create_GPR_model_with_bias(X, Y, mean_function=rhs)

    mu_lhs, var_lhs = model_lhs.predict_f(Xtest)
    mu_rhs, var_rhs = model_rhs.predict_f(Xtest)

    assert_allclose(mu_lhs, mu_rhs)
    assert_allclose(var_lhs, var_rhs)


@pytest.mark.parametrize("mean_functions", [_linear_functions])
def test_linear_mean_functions_associative_property(mean_functions: Sequence[MeanFunction]) -> None:
    """
    Tests that associative property of addition holds for linear mean functions:
    A + (B + (-A)) = B = (A + B) + (-A)
    """
    X, Y = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, Datum.output_dim)
    Xtest = rng.randn(30, Datum.input_dim)
    A, B, A_inverse = mean_functions[0], mean_functions[1], mean_functions[-1]

    lhs = Additive(A, Additive(B, A_inverse))  # A + (B + (-A))
    rhs = Additive(Additive(A, B), A_inverse)  # (A + B) + (-A)

    model_lhs = _create_GPR_model_with_bias(X, Y, mean_function=lhs)
    model_b = _create_GPR_model_with_bias(X, Y, mean_function=B)
    model_rhs = _create_GPR_model_with_bias(X, Y, mean_function=rhs)

    mu_lhs, var_lhs = model_lhs.predict_f(Xtest)
    mu_b, var_b = model_b.predict_f(Xtest)
    mu_rhs, var_rhs = model_rhs.predict_f(Xtest)

    assert_allclose(mu_lhs, mu_b)
    assert_allclose(var_lhs, var_b)
    assert_allclose(mu_b, mu_rhs)
    assert_allclose(var_b, var_rhs)


@pytest.mark.parametrize("batch", [(), (1, 2)])
@pytest.mark.parametrize("degree", [0, 1, 3])
@pytest.mark.parametrize("input_dim", [0, 1, 3])
@pytest.mark.parametrize("output_dim", [1, 2])
def test_polynomial__sanity(
    batch: Tuple[int, ...], degree: int, input_dim: int, output_dim: int
) -> None:
    p = Polynomial(degree, input_dim, output_dim)
    X = np.ones(batch + (input_dim,))
    Y = p(X)
    assert (batch) + (output_dim,) == Y.shape
    assert_allclose(1.0, Y)  # Polynomial is initialised to constant 1.0.


def test_polynomial__compute_powers() -> None:
    assert_allclose(
        [
            (0, 0, 0),
            (0, 0, 1),
            (0, 0, 2),
            (0, 1, 0),
            (0, 1, 1),
            (0, 2, 0),
            (1, 0, 0),
            (1, 0, 1),
            (1, 1, 0),
            (2, 0, 0),
        ],
        list(Polynomial.compute_powers(degree=2, input_dim=3)),
    )


def test_polynomial__1d() -> None:
    # Test on a 1D quadratic function, where we can easily work out the maths.
    p = Polynomial(degree=2, w=[1.0, 2.0, 3.0])
    X = np.array([[1.0], [2.0]])
    Y = p(X)
    assert_allclose(
        [
            [1.0 + 2.0 * 1.0 + 3.0 * (1.0 ** 2)],
            [1.0 + 2.0 * 2.0 + 3.0 * (2.0 ** 2)],
        ],
        Y,
    )


def test_polynomial__linear() -> None:
    # Test on a 3D linear function, where we can easily work out the maths.
    p = Polynomial(degree=1, input_dim=3, w=[1.0, 2.0, 3.0, 4.0])
    X = np.array([1.0, 2.0, 3.0])
    Y = p(X)
    assert_allclose([1.0 + 2.0 * 3.0 + 3.0 * 2.0 + 4.0 * 1.0], Y)


@pytest.mark.parametrize("N, D", [[10, 3]])
def test_switched_mean_function(N: int, D: int) -> None:
    """
    Test for the SwitchedMeanFunction.
    """
    X: AnyNDArray = np.hstack([rng.randn(N, D), 1.0 * rng.randint(0, 2, N).reshape(-1, 1)])
    zeros, ones = Constant(np.zeros(1)), Constant(np.ones(1))
    switched_mean = SwitchedMeanFunction([zeros, ones])

    np_list: AnyNDArray = np.array([0.0, 1.0])
    result_ref = (np_list[X[:, D].astype(default_int())]).reshape(-1, 1)
    result = switched_mean(X)

    assert_allclose(result, result_ref)


def test_bug_277_regression() -> None:
    """
    See github issue #277. This is a regression test.
    """
    model1, model2 = Linear(), Linear()
    assert model1.b.numpy() == model2.b.numpy()
    model2.b.assign([1.0])
    assert not model1.b.numpy() == model2.b.numpy()


_model_classes = [
    gpflow.models.GPR,
    gpflow.models.SGPR,
    gpflow.models.GPRFITC,
    gpflow.models.SVGP,
    gpflow.models.VGP,
    gpflow.models.GPMC,
    gpflow.models.SGPMC,
]


@pytest.mark.parametrize("model_class", _model_classes)
def test_models_with_mean_functions_changes(model_class: Type[Any]) -> None:
    """
    Simply check that all models have a higher prediction with a constant mean
    function than with a zero mean function.

    For compositions of mean functions check that multiplication/ addition of
    a constant results in a higher prediction, whereas addition of zero/
    mutliplication with one does not.
    """
    data = rng.randn(Datum.N, Datum.input_dim), rng.randn(Datum.N, 1)
    Xnew = rng.randn(Datum.Ntest, Datum.input_dim)
    inducing_variable = InducingPoints(Z=rng.randn(Datum.M, Datum.input_dim))
    kernel = gpflow.kernels.Matern32()
    likelihood = gpflow.likelihoods.Gaussian()
    zero_mean = Zero()
    non_zero_mean = Constant(c=np.ones(1) * 10)

    if model_class in [gpflow.models.GPR]:
        model_zero_mean = model_class(data, kernel=kernel, mean_function=zero_mean)
        model_non_zero_mean = model_class(data, kernel=kernel, mean_function=non_zero_mean)
    elif model_class in [gpflow.models.VGP]:
        model_zero_mean = model_class(
            data, likelihood=likelihood, kernel=kernel, mean_function=zero_mean
        )
        model_non_zero_mean = model_class(
            data, likelihood=likelihood, kernel=kernel, mean_function=non_zero_mean
        )
    elif model_class in [gpflow.models.SVGP]:
        model_zero_mean = model_class(
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=zero_mean,
        )
        model_non_zero_mean = model_class(
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=non_zero_mean,
        )
    elif model_class in [gpflow.models.SGPR, gpflow.models.GPRFITC]:
        model_zero_mean = model_class(
            data,
            kernel=kernel,
            inducing_variable=inducing_variable,
            mean_function=zero_mean,
        )
        model_non_zero_mean = model_class(
            data,
            kernel=kernel,
            inducing_variable=inducing_variable,
            mean_function=non_zero_mean,
        )
    elif model_class in [gpflow.models.SGPMC]:
        model_zero_mean = model_class(
            data,
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=zero_mean,
        )
        model_non_zero_mean = model_class(
            data,
            kernel=kernel,
            likelihood=likelihood,
            inducing_variable=inducing_variable,
            mean_function=non_zero_mean,
        )
    elif model_class in [gpflow.models.GPMC]:
        model_zero_mean = model_class(
            data, kernel=kernel, likelihood=likelihood, mean_function=zero_mean
        )
        model_non_zero_mean = model_class(
            data, kernel=kernel, likelihood=likelihood, mean_function=non_zero_mean
        )
    else:
        raise NotImplementedError

    mu_zero, var_zero = model_zero_mean.predict_f(Xnew)
    mu_non_zero, var_non_zero = model_non_zero_mean.predict_f(Xnew)
    # predictive variance remains unchanged after modifying mean function
    assert np.all(var_zero.numpy() == var_non_zero.numpy())
    # predictive mean changes after modifying mean function
    assert not np.all(mu_zero.numpy() == mu_non_zero.numpy())


# tests/gpflow/test_inducing_variables.py

# Copyright 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 typing import Callable, Optional, Tuple

import numpy as np
import pytest
import tensorflow as tf

import gpflow
import gpflow.inducing_variables as giv
from gpflow.experimental.check_shapes import get_shape
from tests.gpflow.experimental.check_shapes.utils import TestContext


def test_inducing_points_with_variable_shape() -> None:
    N, M1, D, P = 50, 13, 3, 1
    X, Y = np.random.randn(N, D), np.random.randn(N, P)

    Z1 = np.random.randn(M1, D)

    # use explicit tf.Variable with None shape:
    iv = giv.InducingPoints(
        tf.Variable(Z1, trainable=False, dtype=gpflow.default_float(), shape=(None, D))
    )
    # Note that we cannot have Z be trainable if we want to be able to change its shape;
    # TensorFlow optimizers expect shape to be known at construction time.

    m = gpflow.models.SGPR(data=(X, Y), kernel=gpflow.kernels.Matern32(), inducing_variable=iv)

    # Check 1: that we can still optimize with None shape
    opt = tf.optimizers.Adam()

    @tf.function
    def optimization_step() -> None:
        opt.minimize(m.training_loss, m.trainable_variables)

    optimization_step()

    # Check 2: that we can successfully assign a new Z with different number of inducing points!
    Z2 = np.random.randn(M1 + 1, D)
    m.inducing_variable.Z.assign(Z2)

    # Check 3: that we can also optimize with changed Z tensor
    optimization_step()


@pytest.mark.parametrize(
    "iv_factory,expected_shape",
    [
        (lambda t: giv.InducingPoints(t), (7, 3, 1)),
        (lambda t: giv.Multiscale(t, t), (7, 3, 1)),
        (lambda t: giv.InducingPatches(t), (7, 3, 1)),
        (
            lambda t: giv.FallbackSharedIndependentInducingVariables(giv.InducingPoints(t)),
            (7, 3, None),
        ),
        (
            lambda t: giv.FallbackSeparateIndependentInducingVariables(
                [
                    giv.InducingPoints(t),
                    giv.InducingPoints(t),
                ]
            ),
            (7, 3, 2),
        ),
        (
            lambda t: giv.SharedIndependentInducingVariables(giv.InducingPoints(t)),
            (7, 3, None),
        ),
        (
            lambda t: giv.SeparateIndependentInducingVariables(
                [
                    giv.InducingPoints(t),
                    giv.InducingPoints(t),
                ]
            ),
            (7, 3, 2),
        ),
    ],
)
def test_shape(
    iv_factory: Callable[[tf.Tensor], giv.InducingVariables],
    expected_shape: Tuple[Optional[int], ...],
) -> None:
    M = expected_shape[0]
    ones = tf.ones((7, 3))
    t1 = tf.Variable(ones)
    t2 = tf.Variable(ones, shape=tf.TensorShape(None))

    for t, e in zip((t1, t2), (expected_shape, None)):
        iv = iv_factory(t)
        assert e == get_shape(iv, TestContext())
        assert M == iv.num_inducing
        assert M == len(iv)


# tests/gpflow/test_initial_value.py

# Copyright 2022 the GPflow authors.
#
# 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.
import ast
from pathlib import Path
from typing import List, Sequence, Union

import pytest

import gpflow


def find_py_files() -> Sequence[Path]:
    """
    Find all ``.py`` files of the GPflow user-facing code.
    """
    root_module = gpflow
    root_path_str = root_module.__file__
    assert root_path_str is not None
    root_path = Path(root_path_str).parent
    return tuple(root_path.glob("**/*.py"))


_PY_FILES = find_py_files()


@pytest.mark.parametrize("path", _PY_FILES, ids=str)
def test_assignment_value(path: Path) -> None:
    """
    Python allows assignments without a value, generally used for typing. For example::

        @tf.function
        def foo(a: tf.Tensor, b: tf.Tensor) -> tf.Tensor:
            result: tf.Tensor  # <-- Assignment without a value. Only for typing.
            result = a + b
            return result

    However, ``TensorFlow < 2.9.0`` cannot compile these.

    This test parses all our ``.py`` files and look for assignments without a value, to ensure all
    our code can be compiled by TensorFlow.

    See also: https://github.com/tensorflow/tensorflow/issues/56024
    """

    tree = ast.parse(path.read_text(), str(path))

    class RequireAssignmentValue(ast.NodeVisitor):
        """
        When ``.visit(tree)`` is called this will recursively visit all nodes of ``tree``.

        Visitor methods are called based on the types of the nodes of the tree. If a node has type
        ``Foo``, then ``visit_Foo(...)`` below will be called. If no such method exists
        ``generic_visit`` is called insted, which just visits the children of the node recursively.
        """

        def __init__(self) -> None:
            super().__init__()

            self.in_function_stack: List[bool] = []
            """
            Keeps track of whether we're inside a function.  We only raise errors for assignments in
            functions - I think variables on classes are handled differently, and doesn't cause
            these problems.
            """

            self.bad_assigns: List[ast.AST] = []
            """
            List of bad assignments, so that we can report multiple errors at the same time.
            """

        def set_in_function(self, node: ast.AST, in_function: bool) -> None:
            """
            Set whether we're inside a function, and visit children recursively.
            """
            self.in_function_stack.append(in_function)
            self.generic_visit(node)
            self.in_function_stack.pop()

        def visit_ClassDef(self, node: ast.ClassDef) -> None:
            """
            Called for class definitions.
            """
            self.set_in_function(node, False)

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            """
            Called for function (and method) definitions.
            """
            self.set_in_function(node, True)

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            """
            Called for async function definitions.
            """
            self.set_in_function(node, True)

        def check_assign(self, node: Union[ast.Assign, ast.AugAssign, ast.AnnAssign]) -> None:
            """
            Checks whether this is a valid assignment, and visit children recursively.
            """
            if self.in_function_stack and self.in_function_stack[-1]:
                if not node.value:
                    self.bad_assigns.append(node)
            self.generic_visit(node)

        def visit_Assign(self, node: ast.Assign) -> None:
            """
            Called for a regular assignment.
            """
            self.check_assign(node)

        def visit_AugAssign(self, node: ast.AugAssign) -> None:
            """
            Called for +=, -=, etc.
            """
            self.check_assign(node)

        def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
            """
            Called for an assigment with a type annotation.
            """
            self.check_assign(node)

    visitor = RequireAssignmentValue()
    visitor.visit(tree)

    if visitor.bad_assigns:
        lines = [""]
        lines.extend(f"{path}:{n.lineno}" for n in visitor.bad_assigns)
        lines.append(
            "TensorFlow cannot compile assignments without a value,"
            " so for performance reasons we require all assgnments to have a value."
            " See: https://github.com/tensorflow/tensorflow/issues/56024 "
        )
        raise AssertionError("\n".join(lines))


# tests/gpflow/test_kullback_leiblers.py

# Copyright 2017 the GPflow authors.
#
# 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.

# -*- coding: utf-8 -*-

from unittest.mock import patch

import numpy as np
import pytest
import tensorflow as tf
from numpy.testing import assert_allclose, assert_almost_equal

import gpflow
from gpflow import Parameter, default_float, default_jitter
from gpflow.base import AnyNDArray, TensorType
from gpflow.experimental.check_shapes import ShapeChecker, check_shapes
from gpflow.inducing_variables import InducingPoints
from gpflow.kernels import Kernel
from gpflow.kullback_leiblers import gauss_kl, prior_kl
from gpflow.utilities.bijectors import triangular

rng = np.random.RandomState(0)

# ------------------------------------------
# Fixtures
# ------------------------------------------

Ln = 2
Nn = 10
Mn = 50


@pytest.fixture(scope="module")
def kernel() -> Kernel:
    k = gpflow.kernels.Matern32() + gpflow.kernels.White()
    k.kernels[1].variance.assign(0.01)
    return k


@pytest.fixture(scope="module")
@check_shapes(
    "return: [N, 1, 1]",
)
def inducing_points() -> InducingPoints:
    return InducingPoints(rng.randn(Nn, 1))


@pytest.fixture(scope="module")
@check_shapes(
    "return: [N, L]",
)
def mu() -> Parameter:
    return Parameter(rng.randn(Nn, Ln))


# ------------------------------------------
# Helpers
# ------------------------------------------


@check_shapes(
    "return: [N, M, M]",
)
def make_sqrt(N: int, M: int) -> TensorType:
    return np.array([np.tril(rng.randn(M, M)) for _ in range(N)])


@check_shapes(
    "return: [N, M, M]",
)
def make_K_batch(N: int, M: int) -> TensorType:
    K_np = rng.randn(N, M, M)
    beye: AnyNDArray = np.array([np.eye(M) for _ in range(N)])
    return 0.1 * (K_np + np.transpose(K_np, (0, 2, 1))) + beye


@check_shapes(
    "q_mu: [broadcast batch...]",
    "q_sigma: [broadcast batch...]",
    "p_var: [broadcast batch...]",
    "return: []",
)
def compute_kl_1d(q_mu: TensorType, q_sigma: TensorType, p_var: TensorType = 1.0) -> TensorType:
    p_var = tf.ones_like(q_sigma) if p_var is None else p_var
    q_var = tf.square(q_sigma)
    kl = 0.5 * (q_var / p_var + tf.square(q_mu) / p_var - 1 + tf.math.log(p_var / q_var))
    return tf.reduce_sum(kl)


# ------------------------------------------
# Data classes: storing constants
# ------------------------------------------


class Datum:
    cs = ShapeChecker().check_shape

    M, N = 5, 4

    mu = cs(rng.randn(M, N), "[M, N]")
    A = cs(rng.randn(M, M), "[M, M]")
    I = cs(np.eye(M), "[M, M]")
    K = cs(A @ A.T + default_jitter() * I, "[M, M]")
    sqrt = cs(make_sqrt(N, M), "[N, M, M]")
    sqrt_diag = cs(rng.randn(M, N), "[M, N]")
    K_batch = cs(make_K_batch(N, M), "[N, M, M]")
    K_cholesky = cs(np.linalg.cholesky(K), "[M, M]")


@pytest.mark.parametrize("diag", [True, False])
def test_kl_k_cholesky(diag: bool) -> None:
    """
    Test that passing K or K_cholesky yield the same answer
    """
    q_mu = Datum.mu
    q_sqrt = Datum.sqrt_diag if diag else Datum.sqrt
    kl_K = gauss_kl(q_mu, q_sqrt, K=Datum.K)
    kl_K_chol = gauss_kl(q_mu, q_sqrt, K_cholesky=Datum.K_cholesky)

    np.testing.assert_allclose(kl_K.numpy(), kl_K_chol.numpy())


@pytest.mark.parametrize("white", [True, False])
def test_diags(white: bool) -> None:
    """
    The covariance of q(x) can be Cholesky matrices or diagonal matrices.
    Here we make sure the behaviours overlap.
    """
    # the chols are diagonal matrices, with the same entries as the diag representation.
    chol_from_diag = tf.stack(
        [tf.linalg.diag(Datum.sqrt_diag[:, i]) for i in range(Datum.N)]  # [N, M, M]
    )
    kl_diag = gauss_kl(Datum.mu, Datum.sqrt_diag, Datum.K if white else None)
    kl_dense = gauss_kl(Datum.mu, chol_from_diag, Datum.K if white else None)

    np.testing.assert_allclose(kl_diag, kl_dense)


@pytest.mark.parametrize("diag", [True, False])
def test_whitened(diag: bool) -> None:
    """
    Check that K=Identity and K=None give same answer
    """
    chol_from_diag = tf.stack(
        [tf.linalg.diag(Datum.sqrt_diag[:, i]) for i in range(Datum.N)]  # [N, M, M]
    )
    s = Datum.sqrt_diag if diag else chol_from_diag

    kl_white = gauss_kl(Datum.mu, s)
    kl_nonwhite = gauss_kl(Datum.mu, s, Datum.I)

    np.testing.assert_allclose(kl_white, kl_nonwhite)


@pytest.mark.parametrize("shared_k", [True, False])
@pytest.mark.parametrize("diag", [True, False])
def test_sumkl_equals_batchkl(shared_k: bool, diag: bool) -> None:
    """
    gauss_kl implicitely performs a sum of KL divergences
    This test checks that doing the sum outside of the function is equivalent
    For q(X)=prod q(x_l) and p(X)=prod p(x_l), check that sum KL(q(x_l)||p(x_l)) = KL(q(X)||p(X))
    Here, q(X) has covariance [L, M, M]
    p(X) has covariance [L, M, M] ( or [M, M] )
    Here, q(x_i) has covariance [1, M, M]
    p(x_i) has covariance [M, M]
    """
    s = Datum.sqrt_diag if diag else Datum.sqrt
    kl_batch = gauss_kl(Datum.mu, s, Datum.K if shared_k else Datum.K_batch)
    kl_sum = []
    for n in range(Datum.N):
        q_mu_n = Datum.mu[:, n][:, None]  # [M, 1]
        q_sqrt_n = (
            Datum.sqrt_diag[:, n][:, None] if diag else Datum.sqrt[n, :, :][None, :, :]
        )  # [1, M, M] or [M, 1]
        K_n = Datum.K if shared_k else Datum.K_batch[n, :, :][None, :, :]  # [1, M, M] or [M, M]
        kl_n = gauss_kl(q_mu_n, q_sqrt_n, K=K_n)
        kl_sum.append(kl_n)

    kl_sum = tf.reduce_sum(kl_sum)
    assert_almost_equal(kl_sum, kl_batch)


@patch("tensorflow.__version__", "2.1.0")
def test_sumkl_equals_batchkl_shared_k_not_diag_mocked_tf21() -> None:
    """
    Version of test_sumkl_equals_batchkl with shared_k=True and diag=False
    that tests the TensorFlow < 2.2 workaround with tiling still works.
    """
    kl_batch = gauss_kl(Datum.mu, Datum.sqrt, Datum.K)
    kl_sum = []
    for n in range(Datum.N):
        q_mu_n = Datum.mu[:, n][:, None]  # [M, 1]
        q_sqrt_n = Datum.sqrt[n, :, :][None, :, :]  # [1, M, M] or [M, 1]
        K_n = Datum.K  # [1, M, M] or [M, M]
        kl_n = gauss_kl(q_mu_n, q_sqrt_n, K=K_n)
        kl_sum.append(kl_n)

    kl_sum = tf.reduce_sum(kl_sum)
    assert_almost_equal(kl_sum, kl_batch)


@pytest.mark.parametrize("dim", [0, 1])
@pytest.mark.parametrize("white", [True, False])
def test_oned(white: bool, dim: bool) -> None:
    """
    Check that the KL divergence matches a 1D by-hand calculation.
    """
    mu1d = Datum.mu[dim, :][None, :]  # [1, N]
    s1d = Datum.sqrt[:, dim, dim][:, None, None]  # [N, 1, 1]
    K1d = Datum.K_batch[:, dim, dim][:, None, None]  # [N, 1, 1]

    kl = gauss_kl(mu1d, s1d, K1d if not white else None)
    kl_1d = compute_kl_1d(
        tf.reshape(mu1d, (-1,)),  # N
        tf.reshape(s1d, (-1,)),  # N
        None if white else tf.reshape(K1d, (-1,)),
    )  # N
    np.testing.assert_allclose(kl, kl_1d)


def test_unknown_size_inputs() -> None:
    """
    Test for #725 and #734. When the shape of the Gaussian's mean had at least
    one unknown parameter, `gauss_kl` would blow up. This happened because
    `tf.size` can only output types `tf.int32` or `tf.int64`.
    """
    mu: AnyNDArray = np.ones([1, 4], dtype=default_float())
    sqrt: AnyNDArray = np.ones([4, 1, 1], dtype=default_float())

    known_shape = gauss_kl(*map(tf.constant, [mu, sqrt]))
    unknown_shape = gauss_kl(mu, sqrt)

    np.testing.assert_allclose(known_shape, unknown_shape)


@pytest.mark.parametrize("white", [True, False])
def test_q_sqrt_constraints(
    inducing_points: InducingPoints, kernel: Kernel, mu: AnyNDArray, white: bool
) -> None:
    """Test that sending in an unconstrained q_sqrt returns the same conditional
    evaluation and gradients. This is important to match the behaviour of the KL, which
    enforces q_sqrt is triangular.
    """

    tril: AnyNDArray = np.tril(rng.randn(Ln, Nn, Nn))

    q_sqrt_constrained = Parameter(tril, transform=triangular())
    q_sqrt_unconstrained = Parameter(tril)

    diff_before_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy()
    assert_allclose(diff_before_gradient_step, 0)

    kls = []
    for q_sqrt in [q_sqrt_constrained, q_sqrt_unconstrained]:

        with tf.GradientTape() as tape:
            kl = prior_kl(inducing_points, kernel, mu, q_sqrt, whiten=white)

        grad = tape.gradient(kl, q_sqrt.unconstrained_variable)
        q_sqrt.unconstrained_variable.assign_sub(grad)
        kls.append(kl)

    diff_kls_before_gradient_step = kls[0] - kls[1]

    assert_allclose(diff_kls_before_gradient_step, 0)

    diff_after_gradient_step = (q_sqrt_constrained - q_sqrt_unconstrained).numpy()
    assert_allclose(diff_after_gradient_step, 0)


# tests/gpflow/test_logdensities.py

# Copyright 2018 the GPflow authors.
#
# 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.

import numpy as np
import pytest
import scipy.stats
from numpy.random import randn
from numpy.testing import assert_allclose
from scipy.stats import multivariate_normal as mvn

from gpflow import logdensities
from gpflow.base import AnyNDArray
from gpflow.utilities import to_default_float

rng = np.random.RandomState(1)


@pytest.mark.parametrize("x, mu, var", [(0.9, 0.5, 1.3)])
def test_gaussian(x: float, mu: float, var: float) -> None:
    gpf = logdensities.gaussian(x, mu, var).numpy()
    sps = scipy.stats.norm(loc=mu, scale=np.sqrt(var)).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, mu, var", [(0.9, 0.5, 1.3)])
def test_lognormal(x: float, mu: float, var: float) -> None:
    gpf = logdensities.lognormal(x, mu, var).numpy()
    sps = scipy.stats.lognorm(s=np.sqrt(var), scale=np.exp(mu)).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize(
    "x, p",
    [
        [1, 0.6],
        [0, 0.6],
    ],
)
def test_bernoulli(x: float, p: float) -> None:
    gpf = logdensities.bernoulli(x, p).numpy()
    sps = scipy.stats.bernoulli.logpmf(k=x, p=p)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize(
    "x, lam",
    [
        [0, 1.3],
        [1, 1.3],
        [2, 1.3],
    ],
)
def test_poisson(x: float, lam: float) -> None:
    gpf = logdensities.poisson(x, lam).numpy()
    sps = scipy.stats.poisson.logpmf(k=x, mu=lam)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, scale", [(0.9, 1.3)])
def test_exponential(x: float, scale: float) -> None:
    gpf = logdensities.exponential(x, scale).numpy()
    sps = scipy.stats.expon(loc=0.0, scale=scale).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, shape, scale", [(0.9, 0.5, 1.3)])
def test_gamma(x: float, shape: float, scale: float) -> None:
    gpf = logdensities.gamma(x, shape, scale).numpy()
    sps = scipy.stats.gamma(a=shape, loc=0.0, scale=scale).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize(
    "x, mean, scale, df",
    [
        (0.9, 0.5, 1.3, 1),
        (0.9, 0.5, 1.3, 2),
        (0.9, 0.5, 1.3, 3),
    ],
)
def test_student_t(x: float, mean: float, scale: float, df: int) -> None:
    cast = to_default_float
    gpf = logdensities.student_t(cast(x), cast(mean), cast(scale), df).numpy()
    sps = scipy.stats.t(df=df, loc=mean, scale=scale).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, alpha, beta", [(0.9, 0.5, 1.3)])
def test_beta(x: float, alpha: float, beta: float) -> None:
    gpf = logdensities.beta(x, alpha, beta).numpy()
    sps = scipy.stats.beta(a=alpha, b=beta).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x, mu, sigma", [(0.9, 0.5, 1.3)])
def test_laplace(x: float, mu: float, sigma: float) -> None:
    gpf = logdensities.laplace(x, mu, sigma).numpy()
    sps = scipy.stats.laplace(loc=mu, scale=sigma).logpdf(x)
    np.testing.assert_allclose(gpf, sps)


@pytest.mark.parametrize("x", [randn(4, 10), randn(4, 1)])
@pytest.mark.parametrize("mu", [randn(4, 10), randn(4, 1)])
@pytest.mark.parametrize("cov_sqrt", [randn(4, 4), np.eye(4)])
def test_multivariate_normal(x: AnyNDArray, mu: AnyNDArray, cov_sqrt: AnyNDArray) -> None:
    cov = np.dot(cov_sqrt, cov_sqrt.T)
    L = np.linalg.cholesky(cov)

    gp_result = logdensities.multivariate_normal(x, mu, L)

    if mu.shape[1] > 1:
        if x.shape[1] > 1:
            sp_result = [mvn.logpdf(x[:, i], mu[:, i], cov) for i in range(mu.shape[1])]
        else:
            sp_result = [mvn.logpdf(x.ravel(), mu[:, i], cov) for i in range(mu.shape[1])]
    else:
        sp_result = mvn.logpdf(x.T, mu.ravel(), cov)
    assert_allclose(gp_result, sp_result)


# tests/gpflow/test_monitor.py

from pathlib import Path
from typing import Any, List, Optional, Sequence, Union

import numpy as np
import pytest
import tensorflow as tf
from _pytest.capture import CaptureFixture
from matplotlib.axes import Axes
from matplotlib.figure import Figure

import gpflow
from gpflow.experimental.check_shapes import check_shape as cs
from gpflow.experimental.check_shapes import check_shapes
from gpflow.models import GPR, GPModel
from gpflow.monitor import (
    ExecuteCallback,
    ImageToTensorBoard,
    ModelToTensorBoard,
    Monitor,
    MonitorTask,
    MonitorTaskGroup,
    ScalarToTensorBoard,
)


class Data:
    num_data = 20
    num_steps = 2


class DummyTask(MonitorTask):
    def run(self, **kwargs: Any) -> None:
        pass


class DummyStepCallback:
    current_step = 0

    @check_shapes(
        "variables[all]: [...]",
        "values[all]: [...]",
    )
    def callback(
        self, step: int, variables: Sequence[tf.Variable], values: Sequence[tf.Tensor]
    ) -> None:
        self.current_step = step


@pytest.fixture
@check_shapes()
def model() -> GPModel:
    data = (
        cs(np.random.randn(Data.num_data, 2), "[N, 2]"),
        cs(np.random.randn(Data.num_data, 2), "[N, 2]"),
    )
    kernel = gpflow.kernels.SquaredExponential(lengthscales=[1.0, 2.0])
    return GPR(data, kernel, noise_variance=0.01)


@pytest.fixture
def monitor(model: GPModel, tmp_path: Path) -> Monitor:
    tmp_path_str = str(tmp_path)

    @check_shapes(
        "return: []",
    )
    def lml_callback() -> tf.Tensor:
        return model.log_marginal_likelihood()

    def print_callback() -> None:
        print("foo")

    return Monitor(
        MonitorTaskGroup(
            [
                ModelToTensorBoard(tmp_path_str, model),
                ScalarToTensorBoard(tmp_path_str, lml_callback, "lml"),
            ],
            period=2,
        ),
        MonitorTaskGroup(ExecuteCallback(print_callback), period=1),
    )


def _get_size_directory(d: Path) -> int:
    """Calculating the size of a directory (in Bytes)."""
    return sum(f.stat().st_size for f in d.glob("**/*"))


# Smoke tests for the individual tasks
# #####################################


def test_ExecuteCallback() -> None:
    def callback() -> None:
        print("ExecuteCallback test")

    task = ExecuteCallback(callback)
    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ImageToTensorBoard(tmp_path: Path) -> None:
    """Smoke test `ImageToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)

    def plotting_cb(fig: Figure, axes: Axes) -> None:
        axes[0, 0].plot(np.random.randn(2), np.random.randn(2))
        axes[1, 0].plot(np.random.randn(2), np.random.randn(2))
        axes[0, 1].plot(np.random.randn(2), np.random.randn(2))
        axes[1, 1].plot(np.random.randn(2), np.random.randn(2))

    fig_kwargs = dict(figsize=(10, 10))
    subplots_kwargs = dict(sharex=True, nrows=2, ncols=2)
    task = ImageToTensorBoard(
        tmp_path_str, plotting_cb, "image", fig_kw=fig_kwargs, subplots_kw=subplots_kwargs
    )

    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ScalarToTensorBoard(tmp_path: Path) -> None:
    """Smoke test `ScalarToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)

    def scalar_cb() -> float:
        return 0.0

    task = ScalarToTensorBoard(tmp_path_str, scalar_cb, "scalar")
    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ScalarToTensorBoard_with_argument(tmp_path: Path) -> None:
    """Smoke test `ScalarToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)

    def scalar_cb(x: Optional[float] = None) -> float:
        assert x is not None
        return 2 * x

    task = ScalarToTensorBoard(tmp_path_str, scalar_cb, "scalar")
    compiled_task = tf.function(task)
    task(0, x=1.0)
    compiled_task(0, x=1.0)


def test_ScalarToTensorBoard_with_wrong_keyword_argument(tmp_path: Path) -> None:
    tmp_path_str = str(tmp_path)

    def scalar_cb(x: Optional[float] = None) -> float:
        assert x is not None
        return 2 * x

    task = ScalarToTensorBoard(tmp_path_str, scalar_cb, "scalar")
    compiled_task = tf.function(task)

    with pytest.raises(TypeError, match=r"got an unexpected keyword argument 'y'"):
        task(0, y=1.0)

    with pytest.raises(TypeError, match=r"got an unexpected keyword argument 'y'"):
        compiled_task(0, y=1.0)


def test_ModelToTensorboard(model: GPModel, tmp_path: Path) -> None:
    """Smoke test `ModelToTensorBoard` in Eager and Compiled mode."""
    tmp_path_str = str(tmp_path)
    task = ModelToTensorBoard(tmp_path_str, model)
    task(0)
    compiled_task = tf.function(task)
    compiled_task(0)


def test_ExecuteCallback_arguments(capsys: CaptureFixture[str]) -> None:
    def cb1(x: Optional[int] = None, **_: Any) -> None:
        assert x is not None
        print(x)

    def cb2(**_: Any) -> None:
        print(2)

    def cb3(y: Optional[int] = None, **_: Any) -> None:
        assert y is not None
        print(y)

    group1 = MonitorTaskGroup([ExecuteCallback(cb1), ExecuteCallback(cb2)])
    group2 = MonitorTaskGroup(ExecuteCallback(cb3))
    monitor = Monitor(group1, group2)
    monitor(0, x=1, y=3)
    out, _ = capsys.readouterr()
    assert out == "1\n2\n3\n"


# Smoke test Monitor and MonitorTaskGroup
# ########################################


def none() -> None:
    return None


@pytest.mark.parametrize(
    "task_or_tasks",
    [
        ExecuteCallback(none),
        [ExecuteCallback(none)],
        [ExecuteCallback(none), ExecuteCallback(none)],
    ],
)
def test_MonitorTaskGroup_and_Monitor(task_or_tasks: Union[MonitorTask, List[MonitorTask]]) -> None:
    group = MonitorTaskGroup(task_or_tasks, period=2)

    # check that the tasks is actually a list (custom setter)
    isinstance(group.tasks, list)

    # Smoke test the __call__
    group(0)
    compiled_group = tf.function(group)
    compiled_group(0)

    # Smoke test the Monitor wrapper
    monitor = Monitor(group)
    monitor(0)
    compiled_monitor = tf.function(monitor)
    compiled_monitor(0)


def test_Monitor(monitor: Monitor) -> None:
    monitor(0)
    compiled_monitor = tf.function(monitor)
    compiled_monitor(0)


# Functionality tests
# ###################


def test_compiled_execute_callable(capsys: CaptureFixture[str]) -> None:
    """
    Test that the `ExecuteCallback` when compiled behaves as expected.
    We test that python prints are not executed anymore.
    """
    string_to_print = "Eager mode"

    def callback() -> None:
        print(string_to_print)

    task = ExecuteCallback(callback)

    # Eager mode
    for i in range(Data.num_steps):
        task(i)
    out, _ = capsys.readouterr()

    # We expect a print for each step
    assert out == (f"{string_to_print}\n" * Data.num_steps)

    # Autograph mode
    compiled_task = tf.function(task)
    for i in tf.range(Data.num_steps):
        compiled_task(i)
    out, _ = capsys.readouterr()
    assert out == f"{string_to_print}\n"


def test_periodicity_group(capsys: CaptureFixture[str]) -> None:
    """Test that groups are called at different periods."""

    task_a = ExecuteCallback(lambda: print("a", end=" "))
    task_b = ExecuteCallback(lambda: print("b", end=" "))
    task_X = ExecuteCallback(lambda: print("X", end=" "))

    group_often = MonitorTaskGroup([task_a, task_b], period=1)
    group_seldom = MonitorTaskGroup([task_X], period=3)
    monitor = Monitor(group_often, group_seldom)
    for i in range(7):
        monitor(i)

    out, _ = capsys.readouterr()
    expected = "a b X a b a b a b X a b a b a b X "
    assert out == expected

    # AutoGraph mode
    compiled_monitor = tf.function(monitor)
    for i in tf.range(7):
        compiled_monitor(i)

    # When using TF's range and compiling the monitoring we only expected the python prints once.
    out, _ = capsys.readouterr()
    assert "a b X"


def test_logdir_created(monitor: Monitor, model: GPModel, tmp_path: Path) -> None:
    """
    Check that TensorFlow summaries are written.
    """
    # check existence
    assert tmp_path.is_dir()
    size_before = _get_size_directory(tmp_path)
    assert size_before > 0

    opt = tf.optimizers.Adam()
    for step in range(Data.num_steps):
        opt.minimize(model.training_loss, model.trainable_variables)
        monitor(step)

    size_after = _get_size_directory(tmp_path)
    assert size_after > size_before


def test_compile_monitor(monitor: Monitor, model: GPModel) -> None:
    opt = tf.optimizers.Adam()

    @tf.function
    def tf_func(step: tf.Tensor) -> None:
        opt.minimize(model.training_loss, model.trainable_variables)
        monitor(step)

    for step in tf.range(100):
        tf_func(step)


def test_scipy_monitor(monitor: Monitor, model: GPModel) -> None:
    opt = gpflow.optimizers.Scipy()

    opt.minimize(model.training_loss, model.trainable_variables, step_callback=monitor)


def test_scipy_monitor_called(model: GPModel) -> None:
    task = DummyTask()
    monitor = Monitor(MonitorTaskGroup(task, period=1))
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables, step_callback=monitor)
    assert task.current_step > 1


def test_scipy_step_callback_called(model: GPModel) -> None:
    dsc = DummyStepCallback()
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables, step_callback=dsc.callback)
    assert dsc.current_step > 1


# tests/integration/test_docs.py

# Copyright 2022 the GPflow authors.
#
# 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.
import subprocess
from pathlib import Path

import pytest

BUILD_DOCS_PATH = Path(__file__).parent.parent.parent / "doc" / "build_docs.py"


@pytest.mark.docs
def test_docs(tmp_path: Path) -> None:
    assert BUILD_DOCS_PATH.is_file()
    build_path = tmp_path / "build"

    try:
        subprocess.run(
            [
                "python",
                str(BUILD_DOCS_PATH),
                "develop",
                str(build_path),
                "--limit-notebooks",
                "--fail-on-warning",
            ]
        ).check_returncode()
    except subprocess.CalledProcessError as e:
        raise AssertionError(
            "Documentation build had errors / warnings."
            " Please fix."
            " Check both any .rst files you may have modified directly,"
            " and docstrings of your Python code."
        ) from e


# tests/integration/test_dynamic_shapes.py

# Copyright 2019 the GPflow authors.
#
# 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.

import numpy as np
import pytest
import tensorflow as tf

import gpflow
from gpflow.base import AnyNDArray
from gpflow.config import default_float
from gpflow.experimental.check_shapes import ShapeChecker

rng = np.random.RandomState(0)


class Datum:
    cs = ShapeChecker().check_shape

    n_inputs = 1
    n_outputs = 2
    n_outputs_c = 1

    X: AnyNDArray = cs(rng.rand(20, n_inputs) * 10, "[N, n_inputs]")
    Y = cs(np.sin(X) + 0.9 * np.cos(X * 1.6) + rng.randn(*X.shape) * 0.8, "[N, n_outputs_c]")
    Y = cs(np.tile(Y, n_outputs), "[N, n_outputs]")  # identical columns
    Xtest: AnyNDArray = cs(rng.rand(10, n_inputs) * 10, "[N_new, n_inputs]")
    data = (X, Y)

    # for classification:
    Yc = cs(Y[:, :n_outputs_c], "[N, n_outputs_c]")
    cdata = (X, Yc)


def test_vgp() -> None:
    X = tf.Variable(
        tf.zeros((1, Datum.n_inputs), dtype=default_float()), shape=(None, None), trainable=False
    )
    Y = tf.Variable(
        tf.zeros((1, Datum.n_outputs), dtype=default_float()), shape=(None, None), trainable=False
    )

    model = gpflow.models.VGP(
        (X, Y),
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.Gaussian(),
        num_latent_gps=Datum.n_outputs,
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -model.elbo()

    model_closure()  # Trigger compilation.

    gpflow.models.vgp.update_vgp_data(model, (Datum.X, Datum.Y))
    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


@pytest.mark.parametrize("whiten", [True, False])
@pytest.mark.parametrize("q_diag", [True, False])
def test_svgp(whiten: bool, q_diag: bool) -> None:
    model = gpflow.models.SVGP(
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.Gaussian(),
        inducing_variable=Datum.X.copy(),
        q_diag=q_diag,
        whiten=whiten,
        mean_function=gpflow.mean_functions.Constant(),
        num_latent_gps=Datum.n_outputs,
    )
    gpflow.set_trainable(model.inducing_variable, False)

    # test with explicitly unknown shapes:
    tensor_spec = tf.TensorSpec(shape=None, dtype=default_float())

    # lambda because: https://github.com/GPflow/GPflow/issues/1929
    elbo_lambda = lambda data: model.elbo(data)

    elbo = tf.function(
        elbo_lambda,
        input_signature=[(tensor_spec, tensor_spec)],
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -elbo(Datum.data)

    model_closure()  # Trigger compilation.

    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


def test_vgp_multiclass() -> None:
    X = tf.Variable(
        tf.zeros((1, Datum.n_inputs), dtype=default_float()), shape=(None, None), trainable=False
    )
    Yc = tf.Variable(
        tf.zeros((1, Datum.n_outputs_c), dtype=default_float()), shape=(None, None), trainable=False
    )

    num_classes = 3
    model = gpflow.models.VGP(
        (X, Yc),
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.MultiClass(num_classes=num_classes),
        num_latent_gps=num_classes,
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -model.elbo()

    model_closure()  # Trigger compilation.

    gpflow.models.vgp.update_vgp_data(model, (Datum.X, Datum.Yc))
    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


def test_svgp_multiclass() -> None:
    num_classes = 3
    model = gpflow.models.SVGP(
        gpflow.kernels.SquaredExponential(),
        gpflow.likelihoods.MultiClass(num_classes=num_classes),
        inducing_variable=Datum.X.copy(),
        num_latent_gps=num_classes,
    )
    gpflow.set_trainable(model.inducing_variable, False)

    # lambda because: https://github.com/GPflow/GPflow/issues/1929
    elbo_lambda = lambda data: model.elbo(data)

    # test with explicitly unknown shapes:
    tensor_spec = tf.TensorSpec(shape=None, dtype=default_float())
    elbo = tf.function(
        elbo_lambda,
        input_signature=[(tensor_spec, tensor_spec)],
    )

    @tf.function
    def model_closure() -> tf.Tensor:
        return -elbo(Datum.cdata)

    model_closure()  # Trigger compilation.

    opt = gpflow.optimizers.Scipy()

    # simply test whether it runs without erroring...:
    opt.minimize(
        model_closure,
        variables=model.trainable_variables,
        options=dict(maxiter=3),
        compile=True,
    )


# tests/integration/test_linear_noise.py

# Copyright 2019 the GPflow authors.
#
# 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 typing import Callable

import numpy as np
import pytest

import gpflow
from gpflow import default_float
from gpflow.base import AnyNDArray, RegressionData
from gpflow.config import default_float
from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.functions import Linear
from gpflow.kernels import Kernel
from gpflow.likelihoods import Gaussian
from gpflow.models import GPR, GPRFITC, SGPR, SVGP, VGP, GPModel, training_loss_closure
from gpflow.models.util import InducingPointsLike


class Datum:
    cs = ShapeChecker().check_shape

    rng = np.random.default_rng(20220630)
    n = 100
    X: AnyNDArray = cs(rng.random((n, 1), dtype=default_float()), "[N, 1]")
    noise_slope = -0.7
    noise_offset = 0.7
    noise = cs(
        (noise_slope * X + noise_offset) * rng.standard_normal((n, 1), dtype=default_float()),
        "[N, 1]",
    )
    Y = cs(np.sin(5 * X) + noise, "[N, 1]")
    data = X, Y


def create_kernel() -> Kernel:
    return gpflow.kernels.RBF(lengthscales=0.2)


def create_inducing() -> InducingPointsLike:
    Z = np.linspace(0.0, 1.0, 10)[:, None]
    iv = gpflow.inducing_variables.InducingPoints(Z)
    gpflow.set_trainable(iv.Z, False)
    return iv


def create_linear_noise() -> Gaussian:
    return Gaussian(scale=Linear())


def gpr(data: RegressionData) -> GPModel:
    return GPR(
        data,
        kernel=create_kernel(),
        likelihood=create_linear_noise(),
    )


def vgp(data: RegressionData) -> GPModel:
    return VGP(
        data,
        kernel=create_kernel(),
        likelihood=create_linear_noise(),
    )


def sgpr(data: RegressionData) -> GPModel:
    return SGPR(
        data,
        kernel=create_kernel(),
        inducing_variable=create_inducing(),
        likelihood=create_linear_noise(),
    )


def gprfitc(data: RegressionData) -> GPModel:
    return GPRFITC(
        data,
        kernel=create_kernel(),
        inducing_variable=create_inducing(),
        likelihood=create_linear_noise(),
    )


def svgp(data: RegressionData) -> GPModel:
    return SVGP(
        kernel=create_kernel(),
        likelihood=create_linear_noise(),
        inducing_variable=create_inducing(),
    )


CREATE_MODELS = (
    gpr,
    vgp,
    sgpr,
    gprfitc,
    svgp,
)


@pytest.mark.parametrize("create_model", CREATE_MODELS)
def test_infer_noise(create_model: Callable[[RegressionData], GPModel]) -> None:
    model = create_model(Datum.data)
    gpflow.optimizers.Scipy().minimize(
        training_loss_closure(model, Datum.data),
        variables=model.trainable_variables,
    )

    noise_scale = model.likelihood.scale
    np.testing.assert_allclose(Datum.noise_slope, noise_scale.A, atol=0.1)
    np.testing.assert_allclose(Datum.noise_offset, noise_scale.b, atol=0.1)


# tests/integration/test_method_equivalence.py

# Copyright 2019 the GPflow authors.
#
# 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 typing import Optional, Tuple

import numpy as np
import pytest
from numpy.testing import assert_allclose

import gpflow
from gpflow.base import AnyNDArray, TensorType
from gpflow.config import default_jitter
from gpflow.experimental.check_shapes import ShapeChecker
from gpflow.kernels import Kernel
from gpflow.likelihoods import Likelihood
from gpflow.mean_functions import Constant
from gpflow.models import GPModel, maximum_log_likelihood_objective, training_loss_closure

rng = np.random.RandomState(0)


class Datum:
    cs = ShapeChecker().check_shape

    X: AnyNDArray = cs(rng.rand(20, 1) * 10, "[N, D]")
    Y = cs(np.sin(X) + 0.9 * np.cos(X * 1.6) + rng.randn(*X.shape) * 0.8, "[N, 1]")
    Y = cs(np.tile(Y, 2), "[N, P]")  # two identical columns
    Xtest: AnyNDArray = cs(rng.rand(10, 1) * 10, "[N_new, D]")
    data = (X, Y)


class DatumVGP:
    N, Ns, DX, DY = 100, 10, 2, 2
    np.random.seed(1)
    X = np.random.randn(N, DX)
    Xs = np.random.randn(Ns, DX)
    Y = np.random.randn(N, DY)
    q_mu = np.random.randn(N, DY)
    q_sqrt = np.random.randn(DY, N, N)
    q_alpha: AnyNDArray = np.random.randn(N, DX)
    q_lambda: AnyNDArray = np.random.randn(N, DX) ** 2
    data = (X, Y)


def _create_full_gp_model() -> gpflow.models.GPR:
    """
    GP Regression
    """
    full_gp_model = gpflow.models.GPR(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        mean_function=gpflow.mean_functions.Constant(),
    )
    return full_gp_model


def _create_approximate_models() -> Tuple[
    gpflow.models.VGP,
    gpflow.models.SVGP,
    gpflow.models.SVGP,
    gpflow.models.SGPR,
    gpflow.models.GPRFITC,
]:
    """
    1) Variational GP (with the likelihood set to Gaussian)
    2) Sparse variational GP (likelihood is Gaussian, inducing points
       at the data)
    3) Sparse variational GP (as above, but with the whitening rotation
       of the inducing variables)
    4) Sparse variational GP Regression (as above, but there the inducing
       variables are 'collapsed' out, as in Titsias 2009)
    5) FITC Sparse GP Regression
    """
    model_1 = gpflow.models.VGP(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Gaussian(),
        mean_function=gpflow.mean_functions.Constant(),
    )

    model_2 = gpflow.models.SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Gaussian(),
        inducing_variable=Datum.X.copy(),
        q_diag=False,
        whiten=False,
        mean_function=gpflow.mean_functions.Constant(),
        num_latent_gps=Datum.Y.shape[1],
    )
    gpflow.set_trainable(model_2.inducing_variable, False)

    model_3 = gpflow.models.SVGP(
        kernel=gpflow.kernels.SquaredExponential(),
        likelihood=gpflow.likelihoods.Gaussian(),
        inducing_variable=Datum.X.copy(),
        q_diag=False,
        whiten=True,
        mean_function=gpflow.mean_functions.Constant(),
        num_latent_gps=Datum.Y.shape[1],
    )
    gpflow.set_trainable(model_3.inducing_variable, False)

    model_4 = gpflow.models.SGPR(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        inducing_variable=Datum.X.copy(),
        mean_function=Constant(),
    )
    gpflow.set_trainable(model_4.inducing_variable, False)

    model_5 = gpflow.models.GPRFITC(
        (Datum.X, Datum.Y),
        kernel=gpflow.kernels.SquaredExponential(),
        inducing_variable=Datum.X.copy(),
        mean_function=Constant(),
    )
    gpflow.set_trainable(model_5.inducing_variable, False)

    return model_1, model_2, model_3, model_4, model_5


def _create_vgp_model(
    kernel: Kernel,
    likelihood: Likelihood,
    q_mu: Optional[TensorType] = None,
    q_sqrt: Optional[TensorType] = None,
) -> gpflow.models.VGP:
    model_vgp = gpflow.models.VGP((DatumVGP.X, DatumVGP.Y), kernel, likelihood)
    if q_mu is not None and q_sqrt is not None:
        model_vgp.q_mu.assign(q_mu)
        model_vgp.q_sqrt.assign(q_sqrt)
    return model_vgp


def _create_vgpao_model(
    kernel: Kernel,
    likelihood: Likelihood,
    q_alpha: Optional[TensorType],
    q_lambda: Optional[TensorType],
) -> gpflow.models.VGPOpperArchambeau:
    model_vgpoa = gpflow.models.VGPOpperArchambeau(
        (DatumVGP.X, DatumVGP.Y), kernel, likelihood, num_latent_gps=DatumVGP.DY
    )
    model_vgpoa.q_alpha.assign(q_alpha)
    model_vgpoa.q_lambda.assign(q_lambda)

    return model_vgpoa


def _create_svgp_model(
    kernel: Kernel,
    likelihood: Likelihood,
    q_mu: Optional[TensorType],
    q_sqrt: Optional[TensorType],
    whiten: bool,
) -> gpflow.models.SVGP:
    model_svgp = gpflow.models.SVGP(
        kernel,
        likelihood,
        DatumVGP.X.copy(),
        whiten=whiten,
        q_diag=False,
        num_latent_gps=DatumVGP.DY,
    )
    model_svgp.q_mu.assign(q_mu)
    model_svgp.q_sqrt.assign(q_sqrt)
    return model_svgp


@pytest.mark.parametrize("approximate_model", _create_approximate_models())
def test_equivalence(approximate_model: GPModel) -> None:
    """
    With a Gaussian likelihood, and inducing points (where appropriate)
    positioned at the data, many of the gpflow methods are equivalent (perhaps
    subject to some optimization).
    """

    def optimize(model: GPModel) -> None:
        opt = gpflow.optimizers.Scipy()
        loss = training_loss_closure(model, Datum.data)
        opt.minimize(loss, model.trainable_variables, options=dict(maxiter=3000))
        if isinstance(model, gpflow.models.SVGP) and not model.whiten:
            # The (S)VGP model in non-whitened representation has significantly
            # worse optimization behaviour. To get the tests to pass, we need
            # to optimize much harder: we set ftol=gtol=0.0 to enforce
            # continued optimization.
            opt.minimize(
                loss, model.trainable_variables, options=dict(maxiter=7000, ftol=0.0, gtol=0.0)
            )

    gpr_model = _create_full_gp_model()
    optimize(gpr_model)
    optimize(approximate_model)

    gpr_likelihood = gpr_model.log_marginal_likelihood()
    approximate_likelihood = maximum_log_likelihood_objective(approximate_model, Datum.data)
    assert_allclose(approximate_likelihood, gpr_likelihood, rtol=1e-6)

    gpr_kernel_ls = gpr_model.kernel.lengthscales.numpy()
    gpr_kernel_var = gpr_model.kernel.variance.numpy()

    approximate_kernel_ls = approximate_model.kernel.lengthscales.numpy()
    approximate_kernel_var = approximate_model.kernel.variance.numpy()

    assert_allclose(gpr_kernel_ls, approximate_kernel_ls, 1e-4)
    assert_allclose(gpr_kernel_var, approximate_kernel_var, 1e-3)

    gpr_mu, gpr_var = gpr_model.predict_y(Datum.Xtest)
    approximate_mu, approximate_var = approximate_model.predict_y(Datum.Xtest)

    assert_allclose(gpr_mu, approximate_mu, 1e-3)
    assert_allclose(gpr_var, approximate_var, 1e-4)


def test_equivalence_vgp_and_svgp() -> None:
    kernel = gpflow.kernels.Matern52()
    likelihood = gpflow.likelihoods.StudentT()

    svgp_model = _create_svgp_model(kernel, likelihood, DatumVGP.q_mu, DatumVGP.q_sqrt, whiten=True)
    vgp_model = _create_vgp_model(kernel, likelihood, DatumVGP.q_mu, DatumVGP.q_sqrt)

    likelihood_svgp = svgp_model.elbo(DatumVGP.data)
    likelihood_vgp = vgp_model.elbo()
    assert_allclose(likelihood_svgp, likelihood_vgp, rtol=1e-2)

    svgp_mu, svgp_var = svgp_model.predict_f(DatumVGP.Xs)
    vgp_mu, vgp_var = vgp_model.predict_f(DatumVGP.Xs)

    assert_allclose(svgp_mu, vgp_mu)
    assert_allclose(svgp_var, vgp_var)


def test_equivalence_vgp_and_opper_archambeau() -> None:
    kernel = gpflow.kernels.Matern52()
    likelihood = gpflow.likelihoods.StudentT()

    vgp_oa_model = _create_vgpao_model(kernel, likelihood, DatumVGP.q_alpha, DatumVGP.q_lambda)

    K = kernel(DatumVGP.X) + np.eye(DatumVGP.N) * default_jitter()
    L = np.linalg.cholesky(K)
    L_inv = np.linalg.inv(L)
    K_inv = np.linalg.inv(K)

    mean = K @ DatumVGP.q_alpha

    prec_dnn = K_inv[None, :, :] + np.array([np.diag(l ** 2) for l in DatumVGP.q_lambda.T])
    var_dnn = np.linalg.inv(prec_dnn)

    svgp_model_unwhitened = _create_svgp_model(
        kernel, likelihood, mean, np.linalg.cholesky(var_dnn), whiten=False
    )

    mean_white_nd = L_inv.dot(mean)
    var_white_dnn = np.einsum("nN,dNM,mM->dnm", L_inv, var_dnn, L_inv)
    q_sqrt_nnd = np.linalg.cholesky(var_white_dnn)

    vgp_model = _create_vgp_model(kernel, likelihood, mean_white_nd, q_sqrt_nnd)

    likelihood_vgp = vgp_model.elbo()
    likelihood_vgp_oa = vgp_oa_model.elbo()
    likelihood_svgp_unwhitened = svgp_model_unwhitened.elbo(DatumVGP.data)

    assert_allclose(likelihood_vgp, likelihood_vgp_oa, rtol=1e-2)
    assert_allclose(likelihood_vgp, likelihood_svgp_unwhitened, rtol=1e-2)

    vgp_oa_mu, vgp_oa_var = vgp_oa_model.predict_f(DatumVGP.Xs)
    svgp_unwhitened_mu, svgp_unwhitened_var = svgp_model_unwhitened.predict_f(DatumVGP.Xs)
    vgp_mu, vgp_var = vgp_model.predict_f(DatumVGP.Xs)

    assert_allclose(vgp_oa_mu, vgp_mu)
    assert_allclose(vgp_oa_var, vgp_var, rtol=1e-4)  # jitter?
    assert_allclose(svgp_unwhitened_mu, vgp_mu)
    assert_allclose(svgp_unwhitened_var, vgp_var, rtol=1e-4)


class DatumUpper:
    cs = ShapeChecker().check_shape

    rng = np.random.default_rng(123)
    X = cs(rng.random((100, 1)), "[N, D]")
    Y = cs(np.sin(1.5 * 2 * np.pi * X) + rng.standard_normal(X.shape) * 0.1 + 5.3, "[N, P]")
    assert Y.mean() > 5.0, "offset ensures a regression test against the bug fixed by PR #1560"
    data = (X, Y)


def test_upper_bound_few_inducing_points() -> None:
    """
    Test for upper bound for regression marginal likelihood
    """
    model_vfe = gpflow.models.SGPR(
        (DatumUpper.X, DatumUpper.Y),
        gpflow.kernels.SquaredExponential(),
        inducing_variable=DatumUpper.X[:10, :].copy(),
        mean_function=Constant(),
        likelihood=gpflow.likelihoods.Gaussian(scale=gpflow.functions.Linear()),
    )
    opt = gpflow.optimizers.Scipy()

    opt.minimize(
        model_vfe.training_loss,
        variables=model_vfe.trainable_variables,
        options=dict(maxiter=500),
    )

    full_gp = gpflow.models.GPR(
        (DatumUpper.X, DatumUpper.Y),
        kernel=model_vfe.kernel,
        mean_function=model_vfe.mean_function,
        likelihood=model_vfe.likelihood,
    )

    lml_upper = model_vfe.upper_bound()
    lml_vfe = model_vfe.elbo()
    lml_full_gp = full_gp.log_marginal_likelihood()

    assert lml_vfe < lml_full_gp
    assert lml_full_gp < lml_upper


# tests/integration/test_notebooks.py

# Copyright 2017 the GPflow authors.
#
# 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.

import itertools
import sys
import traceback
from pathlib import Path
from typing import Sequence, Set

import jupytext
import nbformat
import pytest
from nbconvert.preprocessors import ExecutePreprocessor
from nbconvert.preprocessors.execute import CellExecutionError

_NOTEBOOK_DIR = (Path(__file__) / "../../../doc/sphinx/notebooks").resolve()


def test_notebook_dir_exists() -> None:
    assert _NOTEBOOK_DIR.is_dir()


# To ignore a notebook, add its full base name (including .pct.py or .md
# extension, but without any directory component). NOTE that if there are
# several notebooks in different directories with the same base name, they will
# all get ignored (change the ignoring check to something else in that
# case, if need be!)
IGNORED_NOTEBOOKS: Set[str] = set()


def get_notebooks() -> Sequence[Path]:
    """
    Returns all notebooks in `_nbpath` that are not ignored.
    """

    def notebook_ignored(nb: Path) -> bool:
        return nb.name in IGNORED_NOTEBOOKS

    # recursively traverse the notebook directory in search for ipython notebooks
    all_py_notebooks = _NOTEBOOK_DIR.glob("**/*.pct.py")
    all_md_notebooks = _NOTEBOOK_DIR.glob("**/*.md")
    all_notebooks = itertools.chain(all_md_notebooks, all_py_notebooks)
    notebooks_to_test = [nb for nb in all_notebooks if not notebook_ignored(nb)]
    return notebooks_to_test


def _preproc() -> ExecutePreprocessor:
    pythonkernel = "python" + str(sys.version_info[0])
    return ExecutePreprocessor(timeout=300, kernel_name=pythonkernel, interrupt_on_timeout=True)


def _exec_notebook(notebook_path: Path) -> None:
    with notebook_path.open() as notebook_file:
        nb = jupytext.read(notebook_file, as_version=nbformat.current_nbformat)
        try:
            meta_data = {"path": str(notebook_path.parent)}
            _preproc().preprocess(nb, {"metadata": meta_data})
        except CellExecutionError as cell_error:
            traceback.print_exc(file=sys.stdout)
            msg = "Error executing the notebook {0}. See above for error.\nCell error: {1}"
            pytest.fail(msg.format(str(notebook_path), str(cell_error)))


@pytest.mark.notebooks
@pytest.mark.parametrize("notebook_path", get_notebooks(), ids=str)
def test_notebook(notebook_path: Path) -> None:
    _exec_notebook(notebook_path)


def test_has_notebooks() -> None:
    assert len(get_notebooks()) >= 30, "there are probably some notebooks that were not discovered"


# setup.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# pylint: skip-file

from setuptools import find_packages, setup

##### Dependencies of GPflow

requirements = [
    "deprecated",
    "lark>=1.1.0",
    "multipledispatch>=0.6",
    "numpy",
    "packaging",
    "scipy",
    "setuptools>=41.0.0",  # to satisfy dependency constraints
    "tabulate",
    "tensorflow-probability>=0.12.0",
    "tensorflow>=2.4.0; platform_system!='Darwin' or platform_machine!='arm64'",
    # NOTE: Support of Apple Silicon MacOS platforms is in an experimental mode
    "tensorflow-macos>=2.4.0; platform_system=='Darwin' and platform_machine=='arm64'",
    # NOTE: once we require tensorflow-probability>=0.12, we can remove our custom deepcopy handling
    "typing_extensions",
]


def read_file(filename: str) -> str:
    with open(filename, encoding="utf-8") as f:
        return f.read().strip()


version = read_file("VERSION")
readme_text = read_file("README.md")

packages = find_packages(".", exclude=["tests"])

setup(
    name="gpflow",
    version=version,
    author="James Hensman, Alex Matthews",
    author_email="james.hensman@gmail.com",
    description="Gaussian process methods in TensorFlow",
    long_description=readme_text,
    long_description_content_type="text/markdown",
    license="Apache License 2.0",
    keywords="machine-learning gaussian-processes kernels tensorflow",
    url="https://www.gpflow.org",
    project_urls={
        "Source on GitHub": "https://github.com/GPflow/GPflow",
        "Documentation": "https://gpflow.github.io/GPflow/",
    },
    packages=packages,
    package_data={"": ["*.lark"]},
    include_package_data=True,
    install_requires=requirements,
    extras_require={"ImageToTensorBoard": ["matplotlib"]},
    python_requires=">=3.7",
    classifiers=[
        "Development Status :: 5 - Production/Stable",
        "Environment :: GPU :: NVIDIA CUDA",
        "Intended Audience :: Developers",
        "Intended Audience :: Education",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: Apache Software License",
        "Natural Language :: English",
        "Operating System :: MacOS :: MacOS X",
        "Operating System :: Microsoft :: Windows",
        "Operating System :: POSIX :: Linux",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Typing :: Typed",
    ],
)


# tests/__init__.py

import warnings

import numpy as np
import tensorflow as tf

import gpflow

warnings.filterwarnings("ignore")
gpflow.config.set_default_float(np.float64)


# tests/conftest.py

# Copyright 2022 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 typing import Iterable

import pytest
from _pytest.logging import LogCaptureFixture

from gpflow.experimental.check_shapes.config import (
    DocstringFormat,
    ShapeCheckingState,
    get_enable_check_shapes,
    get_enable_function_call_precompute,
    get_rewrite_docstrings,
    set_enable_check_shapes,
    set_enable_function_call_precompute,
    set_rewrite_docstrings,
)


@pytest.fixture(autouse=True)
def test_auto_graph_compile(caplog: LogCaptureFixture) -> Iterable[None]:
    yield

    for when in ["setup", "call", "teardown"]:
        for record in caplog.get_records(when):
            assert not record.msg.startswith("AutoGraph could not transform"), record.getMessage()


@pytest.fixture(autouse=True)
def enable_shape_checks() -> Iterable[None]:
    # Ensure that:
    # 1: `check_shapes` is enabled when running these tests.
    # 2: If a test manipulates `check_shapes` settings, they are reset after the test.
    # See also: tests/conftest.py
    old_enable = get_enable_check_shapes()
    old_rewrite_docstrings = get_rewrite_docstrings()
    old_function_call_precompute = get_enable_function_call_precompute()
    set_enable_check_shapes(ShapeCheckingState.ENABLED)
    set_rewrite_docstrings(DocstringFormat.SPHINX)
    set_enable_function_call_precompute(True)
    yield
    set_enable_function_call_precompute(old_function_call_precompute)
    set_rewrite_docstrings(old_rewrite_docstrings)
    set_enable_check_shapes(old_enable)


# tests/benchmark/test_grouping.py
# Copyright 2022 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 typing import Collection, Sequence, Tuple
from unittest.mock import Mock

import pandas as pd
import pytest

from benchmark.grouping import GroupingKey as GK
from benchmark.grouping import GroupingSpec, _iter_by, group
from benchmark.metadata import BenchmarkMetadata
from benchmark.plotter_api import Plotter

METRICS_COLUMNS = ["run_id", "dataset", "model", "metric", "value"]
METRICS_DF = pd.DataFrame(
    [
        ("run_1", "ds1", "md1", "mt1", 1.0),
        ("run_2", "ds1", "md1", "mt1", 2.0),
        ("run_1", "ds2", "md1", "mt1", 3.0),
        ("run_1", "ds1", "md2", "mt1", 4.0),
        ("run_1", "ds1", "md1", "mt2", 5.0),
    ],
    columns=METRICS_COLUMNS,
)

PL1 = Mock(Plotter)
PL1.name = "pl1"
PL2 = Mock(Plotter)
PL2.name = "pl2"

BM1 = Mock(BenchmarkMetadata)
BM1.run_id = "run_1"
BM1.git_branch_name = "branch_1"
BM1.user = "user_1"
BM1.timestamp = "20220805_111111.111111"

BM2 = Mock(BenchmarkMetadata)
BM2.run_id = "run_2"
BM2.git_branch_name = "branch_2"
BM2.user = "user_2"
BM2.timestamp = "20220805_222222.222222"


def test_iter_by() -> None:
    assert [
        [],
        [GK.USER],
        [GK.RAM],
        [GK.HOSTNAME],
        [GK.USER, GK.RAM],
        [GK.TIMESTAMP],
        [GK.USER, GK.HOSTNAME],
        [GK.HOSTNAME, GK.RAM],
        [GK.USER, GK.TIMESTAMP],
        [GK.TIMESTAMP, GK.RAM],
        [GK.USER, GK.HOSTNAME, GK.RAM],
        [GK.HOSTNAME, GK.TIMESTAMP],
        [GK.USER, GK.TIMESTAMP, GK.RAM],
        [GK.USER, GK.HOSTNAME, GK.TIMESTAMP],
        [GK.HOSTNAME, GK.TIMESTAMP, GK.RAM],
        [GK.USER, GK.HOSTNAME, GK.TIMESTAMP, GK.RAM],
    ] == list(
        _iter_by(
            [
                GK.USER,
                GK.HOSTNAME,
                GK.TIMESTAMP,
                GK.RAM,
            ]
        )
    )


@pytest.mark.parametrize(
    "spec,expected",
    [
        # Group by no key:
        (
            GroupingSpec(by=[], minimise=False),
            [
                (
                    (),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
            ],
        ),
        # Group by a metadata key:
        (
            GroupingSpec(by=[GK.GIT_BRANCH_NAME], minimise=False),
            [
                (
                    ("branch_1",),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
                (
                    ("branch_2",),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
            ],
        ),
        # Group by a metrics_df key:
        (
            GroupingSpec(by=[GK.DATASET], minimise=False),
            [
                (
                    ("ds1",),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
                (
                    ("ds2",),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1, PL2],
                ),
            ],
        ),
        # Group by a (the) plotter key:
        (
            GroupingSpec(by=[GK.PLOTTER], minimise=False),
            [
                (
                    (),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    (),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
            ],
        ),
        # Group by all the keys:
        (
            GroupingSpec(
                by=[
                    GK.USER,
                    GK.GIT_BRANCH_NAME,
                    GK.TIMESTAMP,
                    GK.DATASET,
                    GK.MODEL,
                    GK.METRIC,
                    GK.PLOTTER,
                ],
                minimise=False,
            ),
            [
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_1", "branch_1", "20220805_111111.111111", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("user_2", "branch_2", "20220805_222222.222222", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("user_2", "branch_2", "20220805_222222.222222", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
            ],
        ),
        # Group by all the keys, and minimise:
        (
            GroupingSpec(
                by=[
                    GK.USER,
                    GK.GIT_BRANCH_NAME,
                    GK.TIMESTAMP,
                    GK.DATASET,
                    GK.MODEL,
                    GK.METRIC,
                    GK.PLOTTER,
                ],
                minimise=True,
            ),
            [
                (
                    ("branch_1", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt1", 1.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_1", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds1", "md1", "mt2"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md1", "mt2", 5.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_1", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds1", "md2", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds1", "md2", "mt1", 4.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_1", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_1", "ds2", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_1", "ds2", "md1", "mt1", 3.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
                (
                    ("branch_2", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL1],
                ),
                (
                    ("branch_2", "ds1", "md1", "mt1"),
                    pd.DataFrame(
                        [
                            ("run_2", "ds1", "md1", "mt1", 2.0),
                        ],
                        columns=METRICS_COLUMNS,
                    ),
                    [PL2],
                ),
            ],
        ),
    ],
)
def test_group(
    spec: GroupingSpec,
    expected: Sequence[Tuple[Tuple[str, ...], pd.DataFrame, Collection[Plotter]]],
) -> None:
    actual = group(METRICS_DF, [PL1, PL2], [BM1, BM2], spec)

    assert len(expected) == len(actual)
    for expected_gr, actual_gr in zip(expected, actual):
        expected_key, expected_df, expected_plotters = expected_gr
        actual_key, actual_df, actual_plotters = actual_gr

        assert expected_key == actual_key
        pd.testing.assert_frame_equal(expected_df, actual_df.reset_index(drop=True))
        assert expected_plotters == list(actual_plotters)
