code: 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())
  
file_path: GPflow/tests/gpflow/kernels/test_broadcasting.py
imports: from typing import Any, Callable, List, Sequence, Tuple, cast
  import numpy
  import pytest
  import tensorflow
  from check_shapes import check_shape
  from check_shapes import check_shapes
  from numpy.testing import assert_allclose
  import gpflow
  import gpflow.ci_utils
  from gpflow import kernels
  from gpflow.base import AnyNDArray, TensorType
  from gpflow.kernels.categorical import Categorical
label_desc: Use image node_id 4 for calling a global function with example usage: test_broadcasting(kernel, batch_shape, batch2_shape, rng) without return types