Remove deprecated PRNGKeyArray.

PiperOrigin-RevId: 623494831 Change-Id: Ie15c1d517fd7b34662d69d05205151d9c7021da4
google-deepmind · Apr 10, 2024 · 7e7de59 · 7e7de59 · Wangchunyu123 · Apr 28, 2024
1 parent 5208f1b
commit 7e7de59
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/alphafold/model/utils.py b/alphafold/model/utils.py
@@ -163,7 +163,7 @@ def inner(key, shape, **kwargs):
     keys = grid_keys(key, shape)
     signature = (
         '()->()'
-        if isinstance(keys, jax.random.PRNGKeyArray)
+        if jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key)
         else '(2)->()'
     )
     return jnp.vectorize(