def is_col_vec(x):
    """Returns if x is a column vector

    Args:
        x (jax.numpy.ndarray): A JAX array.

    Returns:
        True if x is a column vector.
    """
    return x.ndim == 1 or (x.ndim == 2 and x.shape[1] == 1)