Future-proof alphafold prng test for jax custom PRNG

PRNG changes described in jax-ml/jax#9263. After this change, raw key buffers will no longer be accessible via indexing operations.

PiperOrigin-RevId: 530930051
Change-Id: I02b0759c528c393f7ce1109bae83da13cfd68bfe

Author: Jake VanderPlas

alphafold/model/prng_test.py

@@ -29,8 +29,7 @@ def test_key_reuse(self):
 
     raw_key = safe_key.get()
 
-    self.assertNotEqual(raw_key[0], init_key[0])
-    self.assertNotEqual(raw_key[1], init_key[1])
+    self.assertFalse((raw_key == init_key).all())
 
     with self.assertRaises(RuntimeError):
       safe_key.get()