pub_soft/alphafold
3
0
Fork 0
mirror of https://github.com/google-deepmind/alphafold.git synced 2026-06-04 14:58:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert a change that is not compatible with JAX 0.3.25.

Browse Source 
PiperOrigin-RevId: 578486539
Change-Id: Id0e535ebf75916b5179791d093ca714e343ccb69
This commit is contained in:
Augustin Zidek
2023-11-01 05:31:26 -07:00
committed by Copybara-Service
parent f78c589304
commit f715f016d8
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
alphafold/model/utils.py
View File

@@ -163,7 +163,7 @@ def padding_consistent_rng(f):
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(