Only run required rewrites in JAX tests

Several tests ended up not testing the backend Op implementations due to constant folding of inputs.