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