code: def merge_leading_dims(
      x: Optional[tf.Tensor],
      num_dims: int,
  ) -> Optional[tf.Tensor]:
      """Merges leading dimensions of a tensor.
  
      See :func:`split_leading_dim`.
  
      >>> x = tf.ones([3, 2, 1])
      >>> snt.merge_leading_dims(x, num_dims=2)
      <tf.Tensor: ...shape=(6, 1), ...>
  
      If the rank of ``x`` is less than ``num_dims`` it is returned unchanged:
  
      >>> snt.merge_leading_dims(x, 4)
      <tf.Tensor: ...shape=(3, 2, 1), ...>
  
      If ``x`` is not a :tf:`Tensor` or :tf:`Variable` then is is returned
      unchanged:
  
      >>> snt.merge_leading_dims('not a tensor', 1)
      'not a tensor'
  
      Args:
        x: A :tf:`Tensor` to merge.
        num_dims: The number of leading dimensions to merge.
  
      Returns:
        A :tf:`Tensor` with merged leading dimensions or the input unchanged.
      """
      if x is None or not isinstance(x, (tf.Tensor, tf.Variable)):
          return x
  
      # Check if the rank of the input tensor is well-defined.
      if x.shape.dims is None:
          raise ValueError(
              "Can't merge leading dimensions of tensor of unknown rank."
          )
  
      # We can only merge the num_dims leading dimensions if the rank of the given
      # tensor is sufficiently large.
      if num_dims > x.shape.rank:
          return x
  
      static_shape = [maybe_prod(x.shape[:num_dims])] + x.shape[
          num_dims:
      ]
      if (
          static_shape.is_fully_defined()
      ):  # pytype: disable=attribute-error
          return tf.reshape(x, static_shape)
  
      # Shape can't be inferred statically.
      tensor_shape = tf.shape(x)
      leading_dim = tf.reduce_prod(
          tensor_shape[:num_dims], keepdims=True
      )
      other_dims = tensor_shape[num_dims:]
      dynamic_shape = tf.concat([leading_dim, other_dims], axis=0)
      result = tf.reshape(x, dynamic_shape)
      # We lose some static shape information from the above reduce/slice/concat
      # dance, so we explicitly pass it in from what we computed earlier.
      result.set_shape(static_shape)
      return result
  
file_path: sonnet/sonnet/src/batch_apply.py
imports: from typing import Any, Callable, Optional, Sequence, Union
  import numpy
  from sonnet.src import base
  import tensorflow
  import tree
label_desc: Use image node_id 4 for calling a global function with example usage: merge_leading_dims(x, num_dims) and returns: result, x, x, tf