Remove useless jnp.where and add comments for clarity

Remove useless jnp.where and add comments for clarity #3383