af2_complex running error

[pedrovaliente79]

/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/Bio/Data/SCOPData.py:18: BiopythonDeprecationWarning: The 'Bio.
Data.SCOPData' module will be deprecated in a future release of Biopython in favor of 'Bio.Data.PDBData.
warnings.warn(
/scratch/valiente/dl_binder_design/af2_initial_guess/af2_util.py:14: UserWarning: Import of 'rosetta' as a top-level module is deprecated and
may be removed in 2018, import via 'pyrosetta.rosetta'.
from rosetta import *
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
WARNING:tensorflow:From /scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/tf/input_pipeline.py:151: calling map_fn (from t
ensorflow.python.ops.map_fn) with dtype is deprecated and will be removed in a future version.
Instructions for updating:
Use fn_output_signature instead
PyRosetta-4 2023 [Rosetta PyRosetta4.conda.linux.cxx11thread.serialization.CentOS.python39.Release 2023.33+release.9c16e13c3cc4d3ef76e5869e0a
5b44da70cff686 2023-08-18T11:32:04] retrieved from: http://www.pyrosetta.org
(C) Copyright Rosetta Commons Member Institutions. Created in JHU by Sergey Lyskov and PyRosetta Team.
////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////
WARNING! No GPU detected running AF2 on CPU
////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////

Processing struct with tag: design_ppi_0_dldesign_0_cycle1
The distance between residues 95 and 96 is 57.33 A > limit 3.0 A.
I'm going to insert a chainbreak after residue 95
Running model_1_ptm
Traceback (most recent call last):
File "/scratch/valiente/dl_binder_design/af2_initial_guess/predict.py", line 546, in
if args.debug: af2_runner.process_struct(pdb)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/predict.py", line 282, in process_struct
prediction_result = self.model_runner.apply( self.model_runner.params,
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/src/traceback_util.py", line 166, in reraise_with
filtered_traceback
return fun(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/pjit.py", line 253, in cache_miss
outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/pjit.py", line 161, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/api.py", line 324, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/pjit.py", line 491, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/pjit.py", line 969, in _pjit_jaxpr
jaxpr, final_consts, out_type = _create_pjit_jaxpr(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/linear_util.py", line 345, in memoized_fun
ans = call(fun, *args)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/pjit.py", line 922, in _create_pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in t
race_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/src/interpreters/partial_eval.py", line 2177, in t
race_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/transform.py", line 127, in apply_fn
out, state = f.apply(params, {}, *args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/transform.py", line 383, in apply_fn
out = f(*args, **kwargs)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/model.py", line 60, in _forward_fn
return model(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 428, in wrapped
out = f(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 279, in run_interceptors
return bound_method(*args, **kwargs)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/modules.py", line 385, in call
_, prev = hk.while_loop(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/stateful.py", line 621, in while_loop
val, state = jax.lax.while_loop(pure_cond_fun, pure_body_fun, init_val)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/src/traceback_util.py", line 166, in reraise_with
filtered_traceback
return fun(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py", line 1210, in whil
e_loop
init_vals, init_avals, body_jaxpr, in_tree, *rest = _create_jaxpr(init_val)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py", line 1193, in _cre
ate_jaxpr
body_jaxpr, body_consts, body_tree = _initial_style_jaxpr(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py", line 65, in _init
ial_style_jaxpr
jaxpr, consts, out_tree = _initial_style_open_jaxpr(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py", line 59, in _init
ial_style_open_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in t
race_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/src/interpreters/partial_eval.py", line 2177, in t
race_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/stateful.py", line 616, in pure_body_fun
val = body_fun(val)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/modules.py", line 377, in
get_prev(do_call(x[1], recycle_idx=x[0],
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/modules.py", line 339, in do_call
return impl(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 428, in wrapped
out = f(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 279, in run_interceptors
return bound_method(*args, **kwargs)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/modules.py", line 165, in call
representations = evoformer_module(batch0, is_training)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 428, in wrapped
out = f(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 279, in run_interceptors
return bound_method(*args, **kwargs)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/modules.py", line 1771, in call
template_pair_representation = TemplateEmbedding(c.template, gc)(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 428, in wrapped
out = f(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/module.py", line 279, in run_interceptors
return bound_method(*args, **kwargs)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/modules.py", line 2066, in call
template_pair_representation = mapping.sharded_map(map_fn, in_axes=0)(
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/mapping.py", line 141, in mapped_fn
remainder_shape_dtype = hk.eval_shape(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/stateful.py", line 688, in eval_shape
out_shape = jax.eval_shape(stateless_fun, internal_state(), *args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/src/traceback_util.py", line 166, in reraise_with
filtered_traceback
return fun(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/api.py", line 2807, in eval_shape
out = pe.abstract_eval_fun(wrapped_fun.call_wrapped,
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/interpreters/partial_eval.py", line 670, in ab
stract_eval_fun
_, avals_out, _ = trace_to_jaxpr_dynamic(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/interpreters/partial_eval.py", line 2155, in t
race_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/src/interpreters/partial_eval.py", line 2177, in t
race_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/stateful.py", line 684, in stateless_fun
out = fun(*args, **kwargs)
File "/scratch/valiente/dl_binder_design/af2_initial_guess/alphafold/model/mapping.py", line 139, in apply_fun_to_slice
return fun(*input_slice)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/stateful.py", line 576, in mapped_fun
out, state = mapped_pure_fun(args, state)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/src/traceback_util.py", line 166, in reraise_with
filtered_traceback
return fun(*args, **kwargs)
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/api.py", line 1258, in vmap_f
out_flat = batching.batch(
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/linear_util.py", line 188, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/stateful.py", line 568, in pure_fun
state_out = difference(state_in, internal_state())
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/haiku/_src/stateful.py", line 313, in difference
params_after = jax.tree_multimap(functools.partial(if_changed, is_new_param),
File "/scratch/valiente/anaconda3/envs/dl_binder_design/lib/python3.9/site-packages/jax/_src/deprecations.py", line 53, in getattr
raise AttributeError(f"module {module!r} has no attribute {name!r}")
jax._src.traceback_util.UnfilteredStackTrace: AttributeError: module 'jax' has no attribute 'tree_multimap'

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

[nrbennet]

This error is due to having an old version of JAX installed in your environment. I would recommend installing the AF2 environment I have recently added which is easier to install than the previous environment.