Fix minor bugs arising from `jax.array`s in previously unexpected places

Created by: PhilipVinc

cc @attila-i-szabo this should fix your issue. (And also another one I noticed recently with Hilbert spaces)