Update test_raises in samplers

Created by: PhilipVinc

Because Flax 0.3.3 is raising an error before we do, so we must check for different types of errors