def pad_vector(vec, max_size):
    """Pad a vector to a max_size.

    Args:
      vec: a vector to pad.
      max_size: matrix size requested.

    Returns:
      Given V returns [V, 0]
    """
    size = vec.shape[0]
    assert size <= max_size
    if size == max_size:
        return vec
    pad_size = max_size - size
    zs1 = jnp.zeros([pad_size], dtype=vec.dtype)
    return jnp.concatenate([vec, zs1], 0)
Pad a vector to a max_size.

    Args:
      vec: a vector to pad.
      max_size: matrix size requested.

    Returns:
      Given V returns [V, 0]
    
tools/train/scalable_shampoo/distributed_shampoo.py