def conv_out_shape(in_shape, out_fms, p, k, s):
    """
    Gets the output shape [height, width] for a 2D convolution.
    (Assumes square kernel).

    @param in_shape: The input shape [batch, height, width, channels].
    @param out_fms: The number of feature maps in the output.
    @param p: The padding type (either 'SAME' or 'VALID').
    @param k: The side length of the kernel.
    @param s: The stride.

    @return The shape of the output after convolution.
    """
    # convert p to a number
    if p == 'SAME':
        p = k // 2
    elif p == 'VALID':
        p = 0
    else:
        raise ValueError('p must be "SAME" or "VALID".')

    h, w = in_shape[1:3]
    return [in_shape[0],
            int(((h + (2 * p) - k) / s) + 1),
            int(((w + (2 * p) - k) / s) + 1),
            out_fms]