[shape_poly] Improve testing of vmap test harnesses
Previously, we disabled `check_result` (check that the JAX native and JAX with shape polymorphism produce the same result) for test harnesses that are created by vmap on primitive harnesses if the primitive harness has a custom assertion.
Now we enable that checking even for those harnesses, and we use the same custom assertion.
PiperOrigin-RevId: 530547784
