google-deepmind / google-deepmind/alphafold
TypeError: Dtype object is not a valid JAX array type.
- Dominant language
- Python
- Stars
- 14.9k
- Forks
- 2.9k
- PR merge metrics
- No merged PRs in 30d
Description
I'm trying to run the AF2-Multimer model without docker, however, I am not able to generate prediction results due to certain string-like feature objects (e.g. Sequence, domain_name). It looks like it is currently breaking in the `predict()` function in `model.py` here:
`result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)`
I am confused where these object features, which are initially created as objects with `make_sequence_features()` function in `pipeline.py`, are converted into a numerical representation for JAX. See below for the full error output:
```
python3 run_alphafold.py --fasta_paths=multi_sequence_fasta.fasta --output_dir=test_output --use_gpu_relax=True --model_preset=multimer --data_dir='/home/josh/github_repos/alphafold' --no_db_predict=True
I0215 15:49:01.108000 140333781575488 xla_bridge.py:355] Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker:
I0215 15:49:01.108138 140333781575488 xla_bridge.py:355] Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0215 15:49:01.108185 140333781575488 xla_bridge.py:355] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0215 15:49:01.108542 140333781575488 xla_bridge.py:355] Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.
I0215 15:49:01.108612 140333781575488 xla_bridge.py:355] 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.
W0215 15:49:01.108659 140333781575488 xla_bridge.py:362] No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
I0215 15:49:03.409754 140333781575488 run_alphafold.py:497] 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']
I0215 15:49:03.409890 140333781575488 run_alphafold.py:511] Using random seed 133863477447302851 for the data pipeline
I0215 15:49:03.410050 140333781575488 run_alphafold.py:228] Predicting multi_sequence_fasta
I0215 15:49:03.411233 140333781575488 run_alphafold.py:285] Running model model_1_multimer_v3_pred_0 on multi_sequence_fasta
I0215 15:49:03.411429 140333781575488 model.py:167] Running predict with shape(feat) = {'aatype': (215, 21), 'between_segment_residues': (215,), 'domain_name': (1,), 'residue_index': (215,), 'seq_length': (215,), 'sequence': (1,), 'template_all_atom_positions': (1, 215, 37, 3), 'template_all_atom_masks': (1, 215, 37), 'template_sequence': (1,), 'template_aatype': (1, 215, 22), 'template_confidence_scores': (1, 215), 'template_domain_names': (1,), 'template_release_date': (1,)}
Traceback (most recent call last):
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/dispatch.py", line 616, in _valid_jaxtype
xla.abstractify(arg) # faster than core.get_aval
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/interpreters/xla.py", line 267, in abstractify
if aval_fn: return aval_fn(x)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/interpreters/xla.py", line 288, in _make_shaped_array_for_numpy_array
dtypes.check_valid_dtype(dtype)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/dtypes.py", line 613, in check_valid_dtype
raise TypeError(f"Dtype {dtype} is not a valid JAX array "
TypeError: Dtype object is not a valid JAX array type. Only arrays of numeric types are supported by JAX.
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "run_alphafold.py", line 541, in
app.run(main)
File "/home/josh/.local/lib/python3.8/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/home/josh/.local/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "run_alphafold.py", line 516, in main
predict_structure(
File "run_alphafold.py", line 297, in predict_structure
prediction_result = model_runner.predict(processed_feature_dict,
File "/home/josh/github_repos/alphafold/alphafold/model/model.py", line 171, in predict
result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 163, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/api.py", line 676, in cache_miss
dispatch.check_arg(arg)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/dispatch.py", line 623, in check_arg
if not (isinstance(arg, core.Tracer) or _valid_jaxtype(arg)):
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/dispatch.py", line 618, in _valid_jaxtype
return core.valid_jaxtype(arg)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/core.py", line 1243, in valid_jaxtype
concrete_aval(x)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/core.py", line 1258, in concrete_aval
if handler: return handler(x)
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/abstract_arrays.py", line 53, in canonical_concrete_aval
return ConcreteArray(dtypes.canonicalize_dtype(np.result_type(val)), val,
File "/home/josh/miniconda3/envs/af2/lib/python3.8/site-packages/jax/_src/core.py", line 1489, in __init__
assert self.dtype != np.dtype('O'), val
jax._src.traceback_util.UnfilteredStackTrace: AssertionError: [b'hiding_description']
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:
Traceback (most recent call last):
File "run_alphafold.py", line 541, in
app.run(main)
File "/home/josh/.local/lib/python3.8/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/home/josh/.local/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "run_alphafold.py", line 516, in main
predict_structure(
File "run_alphafold.py", line 297, in predict_structure
prediction_result = model_runner.predict(processed_feature_dict,
File "/home/josh/github_repos/alphafold/alphafold/model/model.py", line 171, in predict
result = self.apply(self.params, jax.random.PRNGKey(random_seed), feat)
AssertionError: [b'hiding_description']
```
Contributor guide
Assessment
This issue has not been assessed yet.