Fix places that were passing jnp.array where a shape was expected.#1258