def kinetic_energy(momentum, mass):
    """Compute the kinetic energy of moving particles

    Arguments:
        momentum (array-like of float): momentum of the particles
        mass (array-like of float): mass of the particles
    """
    if momentum.ndim == 3:
        mass = mass[None, :]
    return 0.5 * (momentum / mass[..., None]).sum(axis=(-2, -1))