Function test_broadcasting

tests/gpflow/kernels/test_broadcasting.py:142–311

(
    kernel: kernels.Kernel,
    batch_shape: Tuple[int, ...],
    batch2_shape: Tuple[int, ...],
    rng: np.random.Generator,
)

)
@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: 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]",
        )