Fix custom prng dtype check

Commit

2 years ago

Fix custom prng dtype check Fix related to https://github.com/google/jax/pull/17709 PiperOrigin-RevId: 567426001