def get_density(discrete_velocities):
    density = jnp.sum(discrete_velocities, axis=(- 1))
    return density
def get_density(discrete_velocities):
    density = jnp.sum(discrete_velocities, axis=(- 1))
    return density