Fix metropolis_hastings_jax for machines which only accept 2d input

Created by: PhilipVinc

Previously the sampler was sending vectors to the machine. Now we are always sending batches.