[JAX] Remove unused test utilities that access JAX internals that are about to be removed.
[JAX] Remove unused test utilities that access JAX internals that are about to be removed.

PiperOrigin-RevId: 441522010

google/trax,google/trax