Revert a change that is not compatible with JAX 0.3.25.

PiperOrigin-RevId: 578486539 Change-Id: Id0e535ebf75916b5179791d093ca714e343ccb69
google-deepmind · Nov 1, 2023 · f715f01 · f715f01
1 parent f78c589
commit f715f01
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 jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key)
+        if isinstance(keys, jax.random.PRNGKeyArray)
         else '(2)->()'
     )
     return jnp.vectorize(