docs: def split_leading_dim(tensor, inputs, n_dims=2):
    """Split the first dimension of a tensor.
  
    Args:
      tensor: Tensor to have its first dimension split.
      inputs: Original reference input to look the dimensions of.
      n_dims: Number of dimensions to split.
  
    Returns:
      The input tensor, with its first dimension split.
    """
    input_shape_static = inputs.get_shape()
    input_shape_list = input_shape_static.as_list()
    tensor_shape_static = tensor.get_shape()
    tensor_shape_list = tensor_shape_static.as_list()
    if (input_shape_static.is_fully_defined()
        and tensor_shape_static.is_fully_defined()):
      new_shape = input_shape_list[:n_dims] + tensor_shape_list[1:]
      return tf.reshape(tensor, new_shape)
  
    # Shape can't be inferred statically.
    dims_after_first = tf.shape(tensor)[1:]
    split_sizes = tf.shape(inputs)[:n_dims]
    known_split_sizes = input_shape_list[:n_dims]
    known_dims_after_first = tensor_shape_list[1:]
    output_size = tf.concat([split_sizes, dims_after_first], 0)
    result = tf.reshape(tensor, output_size)
    result.set_shape(known_split_sizes + known_dims_after_first)
    return result