docs: def merge_leading_dims(array_or_tensor, n_dims=2):
    """Merge the first dimensions of a tensor.
  
    Args:
      array_or_tensor: Tensor to have its first dimensions merged. Can also
          be an array or numerical value, which will be converted to a tensor
          for batch application, if needed.
      n_dims: Number of dimensions to merge.
  
    Returns:
      Either the input value converted to a Tensor, with the requested dimensions
      merged, or the unmodified input value if the input has less than `n_dims`
      dimensions.
  
    Raises:
      ValueError: If the rank of `array_or_tensor` is not well-defined.
    """
    tensor = tf.convert_to_tensor(array_or_tensor)
    tensor_shape_static = tensor.get_shape()
  
    # Check if the rank of the input tensor is well-defined.
    if tensor_shape_static.dims is None:
      raise ValueError("Can't merge leading dimensions of tensor of unknown "
                       "rank.")
    tensor_shape_list = tensor_shape_static.as_list()
  
    # We can only merge the n_dims leading dimensions if the rank of the given
    # tensor is sufficiently large.
    if n_dims > len(tensor_shape_list):
      return array_or_tensor
  
    if tensor_shape_static.is_fully_defined():
      new_shape = (
          [np.prod(tensor_shape_list[:n_dims])] + tensor_shape_list[n_dims:])
  
      return tf.reshape(tensor, new_shape)
  
    # Shape can't be inferred statically.
    tensor_shape = tf.shape(tensor)
    new_first_dim = tf.reduce_prod(tensor_shape[:n_dims], keepdims=True)
    other_dims = tensor_shape[n_dims:]
    new_size = tf.concat([new_first_dim, other_dims], 0)
    result = tf.reshape(tensor, new_size)
  
    if all(value is not None for value in tensor_shape_list[:n_dims]):
      merged_leading_size = np.prod(tensor_shape_list[:n_dims])
    else:
      merged_leading_size = None
    # We need to set the result size of this, as otherwise we won't be able to
    # pass to e.g. a Linear. Here we need to know at least the rank of the tensor.
    result.set_shape([merged_leading_size] + tensor_shape_list[n_dims:])
    return result