google-deepmind / google-deepmind/alphafold

jax._src.traceback_util.UnfilteredStackTrace: AttributeError: module 'jax.dtypes' has no attribute 'prng_key'

Open

#849 opened on Oct 31, 2023 13 comments 0 reactions 0 assignees View on GitHub

Dominant language: Python
Stars: 14.9k
Forks: 2.9k
PR merge metrics: No merged PRs in 30d

Description

i am running alphafold v2.3.2 to predict a multimer on 16CPU,241G,4* V100

```
>3JA9_1|Chains A|Proliferating cell nuclear antigen|Homo sapiens (9606)
MFEARLVQGSILKKVLEALKDLINEACWDISSSGVNLQSMDSSHVSLVQLTLRSEGFDTYRCDRNLAMGVNLTSMSKILKCAGNEDIITLRAEDNADTLALVFEAPNQEKVSDYEMKLMDLDVEQLGIPEQEYSCVVKMPSGEFARICRDLSHIGDAVVISCAKDGVKFSASGELGNGNIKLSQTSNVDKEEEAVTIEMNEPVQLTFALRYLNFFTKATPLSSTVTLSMSADVPLVVEYKIADMGHLKYYLAPKIEDEEGS
>3JA9_2|Chains B|Proliferating cell nuclear antigen|Homo sapiens (9606)
MFEARLVQGSILKKVLEALKDLINEACWDISSSGVNLQSMDSSHVSLVQLTLRSEGFDTYRCDRNLAMGVNLTSMSKILKCAGNEDIITLRAEDNADTLALVFEAPNQEKVSDYEMKLMDLDVEQLGIPEQEYSCVVKMPSGEFARICRDLSHIGDAVVISCAKDGVKFSASGELGNGNIKLSQTSNVDKEEEAVTIEMNEPVQLTFALRYLNFFTKATPLSSTVTLSMSADVPLVVEYKIADMGHLKYYLAPKIEDEEGS
>3JA9_3|Chains C|Proliferating cell nuclear antigen|Homo sapiens (9606)
MFEARLVQGSILKKVLEALKDLINEACWDISSSGVNLQSMDSSHVSLVQLTLRSEGFDTYRCDRNLAMGVNLTSMSKILKCAGNEDIITLRAEDNADTLALVFEAPNQEKVSDYEMKLMDLDVEQLGIPEQEYSCVVKMPSGEFARICRDLSHIGDAVVISCAKDGVKFSASGELGNGNIKLSQTSNVDKEEEAVTIEMNEPVQLTFALRYLNFFTKATPLSSTVTLSMSADVPLVVEYKIADMGHLKYYLAPKIEDEEGS
```

command was :
`python3 /shared/alphafold-main/docker/run_docker.py --data_dir=/fsx/soca/dataset --fasta_paths=/data/home/admin/alphafold2/input/rcsb_pdb_3JA9_multimer.fasta --max_template_date=2022-01-01 --model_preset=multimer --db_preset=full_dbs --output_dir=/data/home/admin/alphafold2/output/ `

and i got error, never see this before:

```
The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/app/alphafold/run_alphafold.py", line 570, in
app.run(main)
File "/opt/conda/lib/python3.10/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/opt/conda/lib/python3.10/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/app/alphafold/run_alphafold.py", line 543, in main
predict_structure(
File "/app/alphafold/run_alphafold.py", line 284, in predict_structure
prediction_result = model_runner.predict(processed_feature_dict,
File "/app/alphafold/alphafold/model/model.py", line 167, in predict
result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/transform.py", line 187, in apply_fn
out, state = f.apply(params, None, *args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/transform.py", line 457, in apply_fn
out = f(*args, **kwargs)
File "/app/alphafold/alphafold/model/model.py", line 77, in _forward_fn
return model(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/opt/conda/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 508, in __call__
num_recycles, _, prev, safe_key = hk.while_loop(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 910, in while_loop
val, state = jax.lax.while_loop(pure_cond_fun, pure_body_fun, init_val)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 903, in pure_body_fun
val = body_fun(val)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 486, in recycle_body
ret = apply_network(prev=prev, safe_key=safe_key2)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 449, in apply_network
return impl(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/opt/conda/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 321, in __call__
repr_shape = hk.eval_shape(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 930, in eval_shape
out_shape = jax.eval_shape(stateless_fun, internal_state(), *args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 926, in stateless_fun
out = fun(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 322, in
lambda: embedding_module(batch, is_training))
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/opt/conda/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 633, in __call__
batch = sample_msa(sample_key, batch, c.num_msa)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 274, in sample_msa
index_order = gumbel_argsort_sample_idx(key.get(), logits)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 110, in gumbel_argsort_sample_idx
z = gumbel_noise(key, logits.shape)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 67, in gumbel_noise
uniform_noise = uniform(
File "/app/alphafold/alphafold/model/utils.py", line 166, in inner
if jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key)
AttributeError: module 'jax.dtypes' has no attribute 'prng_key'
```

full logs :

```
[admin@ip-10-247-132-167 ~]$ docker logs -f 39
I1031 11:58:49.800981 140056938311936 templates.py:858] Using precomputed obsolete pdbs /mnt/obsolete_pdbs_path/obsolete.dat.
I1031 11:58:50.871939 140056938311936 xla_bridge.py:353] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker:
I1031 11:58:51.903923 140056938311936 xla_bridge.py:353] Unable to initialize backend 'rocm': NOT_FOUND: Could not find registered platform with name: "rocm". Available platform names are: Interpreter Host CUDA
I1031 11:58:51.904496 140056938311936 xla_bridge.py:353] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
I1031 11:58:51.904650 140056938311936 xla_bridge.py:353] Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
I1031 11:58:58.515359 140056938311936 run_alphafold.py:524] Have 25 models: ['model_1_multimer_v3_pred_0', 'model_1_multimer_v3_pred_1', 'model_1_multimer_v3_pred_2', 'model_1_multimer_v3_pred_3', 'model_1_multimer_v3_pred_4', 'model_2_multimer_v3_pred_0', 'model_2_multimer_v3_pred_1', 'model_2_multimer_v3_pred_2', 'model_2_multimer_v3_pred_3', 'model_2_multimer_v3_pred_4', 'model_3_multimer_v3_pred_0', 'model_3_multimer_v3_pred_1', 'model_3_multimer_v3_pred_2', 'model_3_multimer_v3_pred_3', 'model_3_multimer_v3_pred_4', 'model_4_multimer_v3_pred_0', 'model_4_multimer_v3_pred_1', 'model_4_multimer_v3_pred_2', 'model_4_multimer_v3_pred_3', 'model_4_multimer_v3_pred_4', 'model_5_multimer_v3_pred_0', 'model_5_multimer_v3_pred_1', 'model_5_multimer_v3_pred_2', 'model_5_multimer_v3_pred_3', 'model_5_multimer_v3_pred_4']
I1031 11:58:58.515562 140056938311936 run_alphafold.py:538] Using random seed 247198152922164255 for the data pipeline
I1031 11:58:58.515847 140056938311936 run_alphafold.py:245] Predicting rcsb_pdb_3JA9_multimer
I1031 11:58:58.538447 140056938311936 pipeline_multimer.py:210] Running monomer pipeline on chain A: 3JA9_1|Chains A|Proliferating cell nuclear antigen|Homo sapiens (9606)
I1031 11:58:58.538740 140056938311936 jackhmmer.py:133] Launching subprocess "/usr/bin/jackhmmer -o /dev/null -A /tmp/tmpe3_s_gj6/output.sto --noali --F1 0.0005 --F2 5e-05 --F3 5e-07 --incE 0.0001 -E 0.0001 --cpu 8 -N 1 /tmp/tmp0isaglgt.fasta /mnt/uniref90_database_path/uniref90.fasta"
I1031 11:58:58.603930 140056938311936 utils.py:36] Started Jackhmmer (uniref90.fasta) query
I1031 12:07:14.808301 140056938311936 utils.py:40] Finished Jackhmmer (uniref90.fasta) query in 496.204 seconds
I1031 12:07:15.079364 140056938311936 jackhmmer.py:133] Launching subprocess "/usr/bin/jackhmmer -o /dev/null -A /tmp/tmpcr0k1q7t/output.sto --noali --F1 0.0005 --F2 5e-05 --F3 5e-07 --incE 0.0001 -E 0.0001 --cpu 8 -N 1 /tmp/tmp0isaglgt.fasta /mnt/mgnify_database_path/mgy_clusters_2022_05.fa"
I1031 12:07:15.080246 140056938311936 utils.py:36] Started Jackhmmer (mgy_clusters_2022_05.fa) query
I1031 12:20:35.566447 140056938311936 utils.py:40] Finished Jackhmmer (mgy_clusters_2022_05.fa) query in 800.486 seconds
I1031 12:20:36.874873 140056938311936 hmmbuild.py:121] Launching subprocess ['/usr/bin/hmmbuild', '--hand', '--amino', '/tmp/tmp41ukra03/output.hmm', '/tmp/tmp41ukra03/query.msa']
I1031 12:20:36.941542 140056938311936 utils.py:36] Started hmmbuild query
I1031 12:20:37.812366 140056938311936 hmmbuild.py:128] hmmbuild stdout:
# hmmbuild :: profile HMM construction from multiple sequence alignments
# HMMER 3.1b2 (February 2015); http://hmmer.org/
# Copyright (C) 2015 Howard Hughes Medical Institute.
# Freely distributed under the GNU General Public License (GPLv3).
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# input alignment file: /tmp/tmp41ukra03/query.msa
# output HMM file: /tmp/tmp41ukra03/output.hmm
# input alignment is asserted as: protein
# model architecture construction: hand-specified by RF annotation
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

# idx name nseq alen mlen eff_nseq re/pos description
#---- -------------------- ----- ----- ----- -------- ------ -----------
1 query 4436 1315 261 7.86 0.590

# CPU time: 0.70u 0.02s 00:00:00.72 Elapsed: 00:00:00.72

stderr:

I1031 12:20:37.812544 140056938311936 utils.py:40] Finished hmmbuild query in 0.871 seconds
I1031 12:20:37.814033 140056938311936 hmmsearch.py:103] Launching sub-process ['/usr/bin/hmmsearch', '--noali', '--cpu', '8', '--F1', '0.1', '--F2', '0.1', '--F3', '0.1', '--incE', '100', '-E', '100', '--domE', '100', '--incdomE', '100', '-A', '/tmp/tmpllsqrtfr/output.sto', '/tmp/tmpllsqrtfr/query.hmm', '/mnt/pdb_seqres_database_path/pdb_seqres.txt']
I1031 12:20:37.883789 140056938311936 utils.py:36] Started hmmsearch (pdb_seqres.txt) query
I1031 12:20:44.574548 140056938311936 utils.py:40] Finished hmmsearch (pdb_seqres.txt) query in 6.691 seconds
I1031 12:20:45.725857 140056938311936 hhblits.py:128] Launching subprocess "/usr/bin/hhblits -i /tmp/tmp0isaglgt.fasta -cpu 4 -oa3m /tmp/tmp036aqlum/output.a3m -o /dev/null -n 3 -e 0.001 -maxseq 1000000 -realign_max 100000 -maxfilt 100000 -min_prefilter_hits 1000 -d /mnt/bfd_database_path/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt -d /mnt/uniref30_database_path/UniRef30_2021_03"
I1031 12:20:45.792595 140056938311936 utils.py:36] Started HHblits query
I1031 12:40:35.259527 140056938311936 utils.py:40] Finished HHblits query in 1189.467 seconds
I1031 12:40:35.406030 140056938311936 templates.py:941] Searching for template for: MFEARLVQGSILKKVLEALKDLINEACWDISSSGVNLQSMDSSHVSLVQLTLRSEGFDTYRCDRNLAMGVNLTSMSKILKCAGNEDIITLRAEDNADTLALVFEAPNQEKVSDYEMKLMDLDVEQLGIPEQEYSCVVKMPSGEFARICRDLSHIGDAVVISCAKDGVKFSASGELGNGNIKLSQTSNVDKEEEAVTIEMNEPVQLTFALRYLNFFTKATPLSSTVTLSMSADVPLVVEYKIADMGHLKYYLAPKIEDEEGS
...

I1031 12:41:38.367300 140056938311936 pipeline.py:234] Uniref90 MSA size: 4473 sequences.
I1031 12:41:38.367510 140056938311936 pipeline.py:235] BFD MSA size: 1795 sequences.
I1031 12:41:38.367598 140056938311936 pipeline.py:236] MGnify MSA size: 501 sequences.
I1031 12:41:38.367679 140056938311936 pipeline.py:237] Final (deduplicated) MSA size: 6644 sequences.
I1031 12:41:38.367957 140056938311936 pipeline.py:239] Total number of templates (NB: this can include bad templates and is later filtered to top 4): 20.
I1031 12:41:38.672444 140056938311936 run_alphafold.py:276] Running model model_1_multimer_v3_pred_0 on rcsb_pdb_3JA9_multimer
I1031 12:41:38.673051 140056938311936 model.py:165] Running predict with shape(feat) = {'aatype': (783,), 'residue_index': (783,), 'seq_length': (), 'msa': (2048, 783), 'num_alignments': (), 'template_aatype': (4, 783), 'template_all_atom_mask': (4, 783, 37), 'template_all_atom_positions': (4, 783, 37, 3), 'asym_id': (783,), 'sym_id': (783,), 'entity_id': (783,), 'deletion_matrix': (2048, 783), 'deletion_mean': (783,), 'all_atom_mask': (783, 37), 'all_atom_positions': (783, 37, 3), 'assembly_num_chains': (), 'entity_mask': (783,), 'num_templates': (), 'cluster_bias_mask': (2048,), 'bert_mask': (2048, 783), 'seq_mask': (783,), 'msa_mask': (2048, 783)}
Traceback (most recent call last):
File "/app/alphafold/run_alphafold.py", line 570, in
app.run(main)
File "/opt/conda/lib/python3.10/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/opt/conda/lib/python3.10/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/app/alphafold/run_alphafold.py", line 543, in main
predict_structure(
File "/app/alphafold/run_alphafold.py", line 284, in predict_structure
prediction_result = model_runner.predict(processed_feature_dict,
File "/app/alphafold/alphafold/model/model.py", line 167, in predict
result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/api.py", line 622, in cache_miss
execute = dispatch._xla_call_impl_lazy(fun_, *tracers, **params)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/dispatch.py", line 236, in _xla_call_impl_lazy
return xla_callable(fun, device, backend, name, donated_invars, keep_unused,
File "/opt/conda/lib/python3.10/site-packages/jax/linear_util.py", line 303, in memoized_fun
ans = call(fun, *args)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/dispatch.py", line 359, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "/opt/conda/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/dispatch.py", line 445, in lower_xla_callable
jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
File "/opt/conda/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 2077, in trace_to_jaxpr_final2
jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main, debug_info)
File "/opt/conda/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 2027, in trace_to_subjaxpr_dynamic2
ans = fun.call_wrapped(*in_tracers_)
File "/opt/conda/lib/python3.10/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/transform.py", line 187, in apply_fn
out, state = f.apply(params, None, *args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/transform.py", line 457, in apply_fn
out = f(*args, **kwargs)
File "/app/alphafold/alphafold/model/model.py", line 77, in _forward_fn
return model(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/opt/conda/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 508, in __call__
num_recycles, _, prev, safe_key = hk.while_loop(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 910, in while_loop
val, state = jax.lax.while_loop(pure_cond_fun, pure_body_fun, init_val)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 1111, in while_loop
init_vals, init_avals, body_jaxpr, in_tree, *rest = _create_jaxpr(init_val)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/lax/control_flow/loops.py", line 1094, in _create_jaxpr
body_jaxpr, body_consts, body_tree = _initial_style_jaxpr(
File "/opt/conda/lib/python3.10/site-packages/jax/_src/lax/control_flow/common.py", line 60, in _initial_style_jaxpr
jaxpr, consts, out_tree = _initial_style_open_jaxpr(
File "/opt/conda/lib/python3.10/site-packages/jax/_src/lax/control_flow/common.py", line 54, in _initial_style_open_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 1981, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/opt/conda/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 1998, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/opt/conda/lib/python3.10/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 903, in pure_body_fun
val = body_fun(val)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 486, in recycle_body
ret = apply_network(prev=prev, safe_key=safe_key2)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 449, in apply_network
return impl(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/opt/conda/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 321, in __call__
repr_shape = hk.eval_shape(
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 930, in eval_shape
out_shape = jax.eval_shape(stateless_fun, internal_state(), *args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/jax/_src/api.py", line 3201, in eval_shape
out = pe.abstract_eval_fun(wrapped_fun.call_wrapped,
File "/opt/conda/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 660, in abstract_eval_fun
_, avals_out, _ = trace_to_jaxpr_dynamic(
File "/opt/conda/lib/python3.10/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/opt/conda/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 1981, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/opt/conda/lib/python3.10/site-packages/jax/interpreters/partial_eval.py", line 1998, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/opt/conda/lib/python3.10/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/opt/conda/lib/python3.10/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/stateful.py", line 926, in stateless_fun
out = fun(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 322, in
lambda: embedding_module(batch, is_training))
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/opt/conda/lib/python3.10/contextlib.py", line 79, in inner
return func(*args, **kwds)
File "/opt/conda/lib/python3.10/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 633, in __call__
batch = sample_msa(sample_key, batch, c.num_msa)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 274, in sample_msa
index_order = gumbel_argsort_sample_idx(key.get(), logits)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 110, in gumbel_argsort_sample_idx
z = gumbel_noise(key, logits.shape)
File "/app/alphafold/alphafold/model/modules_multimer.py", line 67, in gumbel_noise
uniform_noise = uniform(
File "/app/alphafold/alphafold/model/utils.py", line 166, in inner
if jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key)
jax._src.traceback_util.UnfilteredStackTrace: AttributeError: module 'jax.dtypes' has no attribute 'prng_key'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

jax._src.traceback_util.UnfilteredStackTrace: AttributeError: module 'jax.dtypes' has no attribute 'prng_key'

Description

Contributor guide

Assessment

Get new issues in your inbox