code: 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])
  
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 5 for calling a global function with example usage: unroll_batches(f, x, value_rank) and returns: tf, f