Folding jobs fail repeatedly (one example of many with the same error)

Expected Behavior

Batch is able to run through all the queries in a CSV file

Current Behavior

Stops running at certain sequences that cause an internal issue.

input which caused failure:

MSPVPSVYDLLLRLVGHQEVLKDPAQMGKAIETLGKQIEEYKPGSFSSLISGKINSFQFTIPSFVNLLVYFLIFYVSLLVVNKTTRVMLTLLKSLATVAILLLIVCMGVLLLLR:MSPVPSVYDLLLRLVGHQEVLKDPAQMGKAIETLGKQIEEYKPGSFSSLISGKINSFQFTIPSFVNLLVYFLIFYVSLLVVNKTTRVMLTLLKSLATVAILLLIVCMGVLLLLR

Steps to Reproduce (for bugs)

CONFIG:
{ "num_queries": 124, "use_templates": true, "num_relax": 0, "msa_mode": "mmseqs2_uniref_env", "model_type": "alphafold2_multimer_v3", "num_models": 5, "num_recycles": 3, "recycle_early_stop_tolerance": null, "num_ensemble": 1, "model_order": [ 1, 2, 3, 5 ], "keep_existing_results": true, "rank_by": "multimer", "max_seq": 512, "max_extra_seq": 1024, "pair_mode": "unpaired_paired", "host_url": "https://api.colabfold.com", "stop_at_score": 100, "random_seed": 0, "num_seeds": 1, "recompile_padding": 10, "commit": "b4c1bc7cf89bc0bd577c5a9d3c1f7bedc1f74152", "use_dropout": true, "use_cluster_profile": true, "use_fuse": true, "use_bfloat16": true, "version": "1.5.1" }

ColabFold Output (for bugs)

2023-02-07 15:41:10,513 Query 1/124: B6K772_SCHJY__B6K772_SCHJY__228aa (length 228)
2023-02-07 15:41:10,514 Sequence 0 found no templates
2023-02-07 15:41:10,518 Padding length to 238
Traceback (most recent call last):
File "/home/ubuntu/colabfold_batch/colabfold-conda/bin/colabfold_batch", line 8, in
sys.exit(main())
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/colabfold/batch.py", line 1778, in main
run(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/colabfold/batch.py", line 1435, in run
results = predict_structure(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/colabfold/batch.py", line 447, in predict_structure
model_runner.predict(input_features, random_seed=seed, prediction_callback=callback)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/model.py", line 218, in predict
result = self.apply(self.params, key, sub_feat)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/api.py", line 622, in cache_miss
execute = dispatch.xla_call_impl_lazy(fun, *tracers, **params)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/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 "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/linear_util.py", line 303, in memoized_fun
ans = call(fun, *args)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/dispatch.py", line 359, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/dispatch.py", line 445, in lower_xla_callable
jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/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 "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 2027, in trace_to_subjaxpr_dynamic2
ans = fun.call_wrapped(*in_tracers)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/transform.py", line 128, in apply_fn
out, state = f.apply(params, {}, *args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/transform.py", line 357, in apply_fn
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/model.py", line 98, in _forward_fn
return model(batch, is_training=is_training)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 456, in call
ret = apply_network(prev=batch.pop("prev"), safe_key=safe_key)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 441, in apply_network
return impl(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 314, in call
repr_shape = hk.eval_shape(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 951, in eval_shape
out_shape = jax.eval_shape(stateless_fun, internal_state(), *args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/api.py", line 3201, in eval_shape
out = pe.abstract_eval_fun(wrapped_fun.call_wrapped,
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 660, in abstract_eval_fun
_, avals_out, _ = trace_to_jaxpr_dynamic(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1981, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1998, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 947, in stateless_fun
out = fun(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 315, in
lambda: embedding_module(batch, is_training))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 679, in call
extra_evoformer_output, safe_key = extra_evoformer_stack(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 265, in wrapped
ret = _LayerStackNoState(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 156, in call
carry, zs = hk.scan(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 640, in scan
(carry, state), ys = jax.lax.scan(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/control_flow/loops.py", line 259, in scan
init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/control_flow/loops.py", line 245, in _create_jaxpr
jaxpr, consts, out_tree = _initial_style_jaxpr(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/control_flow/common.py", line 60, in _initial_style_jaxpr
jaxpr, consts, out_tree = _initial_style_open_jaxpr(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/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 "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1981, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1998, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 623, in stateful_fun
carry, out = f(carry, x)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 149, in layer
out_x, z = self._call_wrapped(carry.x, *scanned.args_ys)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 182, in _call_wrapped
ret = self._f(*args)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 665, in extra_evoformer_fn
extra_evoformer_output = extra_evoformer_iteration(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules.py", line 1734, in call
msa_act = dropout_wrapper_fn(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules.py", line 76, in dropout_wrapper
residual = module(input_act, mask, is_training=is_training, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules.py", line 820, in call
msa_act = mapping.inference_subbatch(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 223, in inference_subbatch
return sharded_module(*batched_args)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 145, in mapped_fn
remainder_shape_dtype = hk.eval_shape(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 951, in eval_shape
out_shape = jax.eval_shape(stateless_fun, internal_state(), *args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/api.py", line 3201, in eval_shape
out = pe.abstract_eval_fun(wrapped_fun.call_wrapped,
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 660, in abstract_eval_fun
_, avals_out, _ = trace_to_jaxpr_dynamic(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1981, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1998, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/linear_util.py", line 167, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 947, in stateless_fun
out = fun(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 140, in apply_fun_to_slice
input_slice = jax.tree_map(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/tree_util.py", line 207, in tree_map
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/tree_util.py", line 207, in
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 141, in
lambda array, axis: _maybe_slice(array, slice_start, slice_size, axis
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 38, in _maybe_slice
return jax.lax.dynamic_slice_in_dim(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/slicing.py", line 695, in dynamic_slice_in_dim
return dynamic_slice(operand, start_indices, slice_sizes)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/slicing.py", line 109, in dynamic_slice
return dynamic_slice_p.bind(operand, *start_indices, *dynamic_sizes,
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/core.py", line 329, in bind
return self.bind_with_trace(find_top_trace(args), args, params)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/core.py", line 332, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1712, in process_primitive
return custom_staging_rules[primitive](self, *tracers, **params)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/slicing.py", line 911, in dynamic_slice_staging_rule
return trace.default_process_primitive(dynamic_slice_p, (x, *start_indices),
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1717, in default_process_primitive
out_avals, effects = primitive.abstract_eval(*avals, **params)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/core.py", line 365, in abstract_eval
return abstract_eval(*args, **kwargs), no_effects
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/utils.py", line 66, in standard_abstract_eval
return core.ShapedArray(shape_rule(*avals, **kwargs),
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/jax/_src/lax/slicing.py", line 840, in _dynamic_slice_shape_rule
raise TypeError(msg.format(slice_sizes, operand.shape))
jax._src.traceback_util.UnfilteredStackTrace: TypeError: slice slice_sizes must be less than or equal to operand shape, got slice_sizes (4, 228, 64) for operand shape (0, 228, 64).

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 "/home/ubuntu/colabfold_batch/colabfold-conda/bin/colabfold_batch", line 8, in
sys.exit(main())
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/colabfold/batch.py", line 1778, in main
run(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/colabfold/batch.py", line 1435, in run
results = predict_structure(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/colabfold/batch.py", line 447, in predict_structure
model_runner.predict(input_features, random_seed=seed, prediction_callback=callback)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/model.py", line 218, in predict
result = self.apply(self.params, key, sub_feat)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/transform.py", line 128, in apply_fn
out, state = f.apply(params, {}, *args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/transform.py", line 357, in apply_fn
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/model.py", line 98, in _forward_fn
return model(batch, is_training=is_training)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 456, in call
ret = apply_network(prev=batch.pop("prev"), safe_key=safe_key)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 441, in apply_network
return impl(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 314, in call
repr_shape = hk.eval_shape(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 951, in eval_shape
out_shape = jax.eval_shape(stateless_fun, internal_state(), *args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 947, in stateless_fun
out = fun(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 315, in
lambda: embedding_module(batch, is_training))
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 679, in call
extra_evoformer_output, safe_key = extra_evoformer_stack(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 265, in wrapped
ret = _LayerStackNoState(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 156, in call
carry, zs = hk.scan(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 640, in scan
(carry, state), ys = jax.lax.scan(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 623, in stateful_fun
carry, out = f(carry, x)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 149, in layer
out_x, z = self._call_wrapped(carry.x, *scanned.args_ys)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/layer_stack.py", line 182, in _call_wrapped
ret = self._f(*args)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules_multimer.py", line 665, in extra_evoformer_fn
extra_evoformer_output = extra_evoformer_iteration(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules.py", line 1734, in call
msa_act = dropout_wrapper_fn(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules.py", line 76, in dropout_wrapper
residual = module(input_act, mask, is_training=is_training, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 426, in wrapped
out = f(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/module.py", line 272, in run_interceptors
return bound_method(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/modules.py", line 820, in call
msa_act = mapping.inference_subbatch(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 223, in inference_subbatch
return sharded_module(*batched_args)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 145, in mapped_fn
remainder_shape_dtype = hk.eval_shape(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 951, in eval_shape
out_shape = jax.eval_shape(stateless_fun, internal_state(), *args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/haiku/_src/stateful.py", line 947, in stateless_fun
out = fun(*args, **kwargs)
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 140, in apply_fun_to_slice
input_slice = jax.tree_map(
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 141, in
lambda array, axis: _maybe_slice(array, slice_start, slice_size, axis
File "/home/ubuntu/colabfold_batch/colabfold-conda/lib/python3.8/site-packages/alphafold/model/mapping.py", line 38, in _maybe_slice
return jax.lax.dynamic_slice_in_dim(
TypeError: slice slice_sizes must be less than or equal to operand shape, got slice_sizes (4, 228, 64) for operand shape (0, 228, 64).

CONTEXT

Running a series of jobs using a CSV of ids and sequences as input.
Running on Ubuntu 20.04.5 LTS server with an A100 40GB GPU

Start in colabfold/batch.py at run() and predict_structure(), then follow the supplied traceback into alphafold/model/modules_multimer.py and layer_stack.py. Reproduce with the provided configuration and failing CSV sequence. Done means the batch processes all queries instead of stopping on this internal error.