def pad(tensor, pad_len):
    """Pad tensor on first dimension to pad_len.

    Args:
      tensor: input tensor of shape length >= 2
      pad_len: pad length

    Returns:
      tf.Tensor: Padded input tensor.
    """

    assert len(tensor.shape) >= 2  # tensor of shape [batch, length, ...]
    length = tf.shape(tensor)[1]

    padding = [[0, 0], [0, pad_len - length]]
    padding += [[0, 0]] * (len(tensor.shape) - 2)
    return tf.pad(tensor, padding)
Pad tensor on first dimension to pad_len.

    Args:
      tensor: input tensor of shape length >= 2
      pad_len: pad length

    Returns:
      tf.Tensor: Padded input tensor.
    
tensor2tensor/models/research/shuffle_network.py