Error running alphafold

[pachterlab]

Error running alphafold

[xinyangbing]

Can confirm short test sequence finished running without a problem: gget alphafold MAAHKGAEHHHKAAEHHEQAAKHHHAAAEHHEKGEHEQAAHHADTAYAHHKHAEEHAAQAAKHDAEHHAPKPH
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/data_structures.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
PyTreeDef = type(jax.tree_structure(None))
Thu Aug 18 14:25:21 2022 INFO Validating input sequence(s).
Using the single-chain model.
Thu Aug 18 14:25:21 2022 INFO Finding closest source for reference database.
Jackhmmer search: 100%|██████████████████████████████████████████████████████████| 147/147 [elapsed: 1:02:57 remaining: 00:00]
Thu Aug 18 15:30:35 2022 INFO 58 unique sequences found in uniref90 for sequence 1.
Thu Aug 18 15:30:35 2022 INFO 110 unique sequences found in smallbfd for sequence 1.
Thu Aug 18 15:30:35 2022 INFO 9 unique sequences found in mgnify for sequence 1.

173 unique sequences found in total for sequence 1

Running model_1: 0%| | 0/7 [elapsed: 00:00 remaining: ?]Thu Aug 18 15:30:37 2022 WARNING No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/data_structures.py:144: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.
leaves, treedef = jax.tree_flatten(tree)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/data_structures.py:145: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead.
return jax.tree_unflatten(treedef, leaves)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/stateful.py:32: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(lambda x: x, bundle)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/modules.py:330: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
ensembled_batch = jax.tree_map(slice_recycle_idx, batch)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/modules.py:335: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
non_ensembled_batch = jax.tree_map(lambda x: x, prev)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/src/stateful.py:771: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
if not jax.tree_leaves(in_axes):
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:50: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.
values_tree_def = jax.tree_flatten(values)[1]
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:54: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead.
return jax.tree_unflatten(values_tree_def, flat_axes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:128: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
in_sizes = jax.tree_map(maybe_get_size, args, in_axes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:129: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.
flat_sizes = jax.tree_flatten(in_sizes)[0]
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:140: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
input_slice = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:330: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
rotation = jax.tree_map(expand_fn, rotation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:331: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
translation = jax.tree_map(expand_fn, translation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/prng.py:48: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(SafeKey, tuple(new_keys))
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:147: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
out_dtypes = jax.tree_map(lambda x: x.dtype, remainder_shape_dtype)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:148: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
out_shapes = jax.tree_map(lambda x: x.shape, remainder_shape_dtype)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:154: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
shard_shapes = jax.tree_map(lambda x: x.shape, regular_shard_shape_dtype)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:161: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
out_shapes = jax.tree_map(make_output_shape, out_axes, shard_shapes,
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:184: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
outputs = jax.tree_map(allocate_buffer, out_dtypes, out_shapes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/src/stateful.py:575: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
length = jax.tree_leaves(xs)[0].shape[0]
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:173: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(update_slice, outputs, slice_out, out_axes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/stateful.py:314: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
params_after = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/stateful.py:322: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
state_after = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:304: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
rotation = jax.tree_map(expand_fn, rotation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:305: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
translation = jax.tree_map(expand_fn, translation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:499: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi2_frame_to_frame = jax.tree_map(lambda x: x[:, 5], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:500: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi3_frame_to_frame = jax.tree_map(lambda x: x[:, 6], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:501: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi4_frame_to_frame = jax.tree_map(lambda x: x[:, 7], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:503: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi1_frame_to_backb = jax.tree_map(lambda x: x[:, 4], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:516: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
all_frames_to_backb = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:526: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
jax.tree_map(lambda x: x[:, None], backb_to_global),
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:554: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
map_atoms_to_global = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:570: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
pred_positions = jax.tree_map(lambda x: x * mask, pred_positions)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/folding.py:457: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
output = jax.tree_map(lambda *x: jnp.stack(x), *outputs)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/modules.py:319: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(jax.lax.stop_gradient, new_prev)
2022-08-18 15:37:35.147327: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65]

---

[Compiling module jit_apply_fn.5] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

2022-08-18 15:42:51.437904: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 7m16.29072564s

---

[Compiling module jit_apply_fn.5] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/model.py:172: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
jax.tree_map(lambda x: x.block_until_ready(), result)
Running model_2: 14%|█████████ | 1/7 [elapsed: 24:58 remaining: 2:29:51]2022-08-18 16:08:09.881534: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 7m29.936369231s

---

[Compiling module jit_apply_fn.7] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

Running model_3: 29%|██████████████████ | 2/7 [elapsed: 49:08 remaining: 2:02:28]2022-08-18 16:29:46.228967: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 6m40.557215143s

---

[Compiling module jit_apply_fn.8] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

Running model_4: 43%|██████████████████████████▏ | 3/7 [elapsed: 1:12:02 remaining: 1:35:05]2022-08-18 16:47:57.276966: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65]

---

[Compiling module jit_apply_fn.9] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

2022-08-18 16:53:06.778986: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 7m9.502119733s

---

[Compiling module jit_apply_fn.9] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

Running model_5: 57%|██████████████████████████████████▊ | 4/7 [elapsed: 1:35:56 remaining: 1:11:28]2022-08-18 17:17:07.760154: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 7m11.676369326s

---

[Compiling module jit_apply_fn.10] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

Running model_2_ptm: 71%|██████████████████████████████████████████▏ | 5/7 [elapsed: 1:57:52 remaining: 46:16]2022-08-18 17:35:49.388007: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65]

---

[Compiling module jit_apply_fn.11] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

2022-08-18 17:42:01.516711: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 8m12.12882662s

---

[Compiling module jit_apply_fn.11] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

Running model_2_ptm: 86%|██████████████████████████████████████████████████▌ | 6/7 [elapsed: 2:23:14 remaining: 23:53]Thu Aug 18 17:53:49 2022 WARNING
Running model without relaxation stage. Use flag [--relax] ('relax=True') to include AMBER relaxation.
Running model_2_ptm: 100%|███████████████████████████████████████████████████████████| 7/7 [elapsed: 2:23:23 remaining: 00:00]

[lauraluebbert]

Hi, I am having trouble reproducing this error. Jackhmmer will create a "tmp" folder in your home directory where it will temporarily store downloaded database chunks and then delete them. It looks like one of the database chunks was already deleted before the 'os.remove()' command. I am not sure how this happened. Does it happen every time you run alphafold?

[xinyangbing]

Yes this repeatedly happens. Could it be a property of the VM?

[xinyangbing]

Do you think it could be because the AA input is too big? I am trying your test sequence from example, and it is running, but very slowly:
gget alphafold MAAHKGAEHHHKAAEHHEQAAKHHHAAAEHHEKGEHEQAAHHADTAYAHHKHAEEHAAQAAKHDAEHHAPKPH
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/data_structures.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
PyTreeDef = type(jax.tree_structure(None))
Thu Aug 18 14:25:21 2022 INFO Validating input sequence(s).
Using the single-chain model.
Thu Aug 18 14:25:21 2022 INFO Finding closest source for reference database.
Jackhmmer search: 100%|██████████████████████████████████████████████████████████| 147/147 [elapsed: 1:02:57 remaining: 00:00]
Thu Aug 18 15:30:35 2022 INFO 58 unique sequences found in uniref90 for sequence 1.
Thu Aug 18 15:30:35 2022 INFO 110 unique sequences found in smallbfd for sequence 1.
Thu Aug 18 15:30:35 2022 INFO 9 unique sequences found in mgnify for sequence 1.

173 unique sequences found in total for sequence 1

Running model_1: 0%| | 0/7 [elapsed: 00:00 remaining: ?]Thu Aug 18 15:30:37 2022 WARNING No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/data_structures.py:144: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.
leaves, treedef = jax.tree_flatten(tree)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/data_structures.py:145: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead.
return jax.tree_unflatten(treedef, leaves)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/stateful.py:32: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(lambda x: x, bundle)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/modules.py:330: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
ensembled_batch = jax.tree_map(slice_recycle_idx, batch)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/modules.py:335: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
non_ensembled_batch = jax.tree_map(lambda x: x, prev)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/src/stateful.py:771: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
if not jax.tree_leaves(in_axes):
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:50: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.
values_tree_def = jax.tree_flatten(values)[1]
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:54: FutureWarning: jax.tree_unflatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_unflatten instead.
return jax.tree_unflatten(values_tree_def, flat_axes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:128: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
in_sizes = jax.tree_map(maybe_get_size, args, in_axes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:129: FutureWarning: jax.tree_flatten is deprecated, and will be removed in a future release. Use jax.tree_util.tree_flatten instead.
flat_sizes = jax.tree_flatten(in_sizes)[0]
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:140: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
input_slice = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:330: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
rotation = jax.tree_map(expand_fn, rotation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:331: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
translation = jax.tree_map(expand_fn, translation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/prng.py:48: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(SafeKey, tuple(new_keys))
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:147: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
out_dtypes = jax.tree_map(lambda x: x.dtype, remainder_shape_dtype)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:148: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
out_shapes = jax.tree_map(lambda x: x.shape, remainder_shape_dtype)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:154: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
shard_shapes = jax.tree_map(lambda x: x.shape, regular_shard_shape_dtype)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:161: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
out_shapes = jax.tree_map(make_output_shape, out_axes, shard_shapes,
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:184: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
outputs = jax.tree_map(allocate_buffer, out_dtypes, out_shapes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/src/stateful.py:575: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
length = jax.tree_leaves(xs)[0].shape[0]
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/mapping.py:173: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(update_slice, outputs, slice_out, out_axes)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/stateful.py:314: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
params_after = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/haiku/_src/stateful.py:322: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
state_after = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:304: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
rotation = jax.tree_map(expand_fn, rotation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/quat_affine.py:305: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
translation = jax.tree_map(expand_fn, translation)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:499: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi2_frame_to_frame = jax.tree_map(lambda x: x[:, 5], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:500: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi3_frame_to_frame = jax.tree_map(lambda x: x[:, 6], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:501: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi4_frame_to_frame = jax.tree_map(lambda x: x[:, 7], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:503: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
chi1_frame_to_backb = jax.tree_map(lambda x: x[:, 4], all_frames)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:516: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
all_frames_to_backb = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:526: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
jax.tree_map(lambda x: x[:, None], backb_to_global),
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:554: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
map_atoms_to_global = jax.tree_map(
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/all_atom.py:570: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
pred_positions = jax.tree_map(lambda x: x * mask, pred_positions)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/folding.py:457: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
output = jax.tree_map(lambda *x: jnp.stack(x), *outputs)
/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/modules.py:319: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(jax.lax.stop_gradient, new_prev)
2022-08-18 15:37:35.147327: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65]

---

[Compiling module jit_apply_fn.5] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

2022-08-18 15:42:51.437904: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 7m16.29072564s

---

[Compiling module jit_apply_fn.5] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

/home/ccadmin/anaconda3/envs/gget/lib/python3.9/site-packages/alphafold/model/model.py:172: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
jax.tree_map(lambda x: x.block_until_ready(), result)
Running model_2: 14%|█████████ | 1/7 [elapsed: 24:58 remaining: 2:29:51]2022-08-18 16:08:09.881534: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 7m29.936369231s

---

[Compiling module jit_apply_fn.7] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

---

Running model_3: 29%|██████████████████ | 2/7 [elapsed: 49:08 remaining: 2:02:28]

[lauraluebbert]

Could you please send me the sequence that caused a problem before?