bugfix: removing nans from alphafold-multimer+templates

google-deepmind#513

alphafold/model/utils.py

@@ -34,7 +34,7 @@ def final_init(config):
 
 def batched_gather(params, indices, axis=0, batch_dims=0):
  """Implements a JAX equivalent of `tf.gather` with `axis` and `batch_dims`."""
- take_fn = lambda p, i: jnp.take(p, i, axis=axis)
+ take_fn = lambda p, i: jnp.take(p, i, axis=axis, mode="clip")
  for _ in range(batch_dims):
  take_fn = jax.vmap(take_fn)
  return take_fn(params, indices)