bytedance / bytedance/Protenix

running in TFG mode with only RNA causes a crash

Open Beginner friendly
#322 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.1k
Forks
312
PR merge metrics
No merged PRs in 30d

Description

I'm running protenix on some simple RNA loops. When I run without any protein sequence, it crashes. When I add a simple dummy oligopeptide it works again. Here is the error.

```
[Rank 0] rna_loops_ptx_tfg_guidance_AAAA failed: VinaStericPotential._get_collision_candidates.._cache_and_return() takes 1 positional argument but 2 were given
Traceback (most recent call last):
File "/opt/conda/lib/python3.12/site-packages/runner/inference.py", line 486, in infer_predict
prediction = runner.predict(data)
^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/runner/inference.py", line 228, in predict
prediction, _, _ = self.model(
^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/model/protenix.py", line 904, in forward
pred_dict, log_dict, time_tracker = self.main_inference_loop(
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/model/protenix.py", line 432, in main_inference_loop
return self._main_inference_loop(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/model/protenix.py", line 562, in _main_inference_loop
pred_dict["coordinate"] = self.sample_diffusion(
^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/model/protenix.py", line 341, in sample_diffusion
return autocasting_disable_decorator(self.configs.skip_amp.sample_diffusion)(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/utils/torch_utils.py", line 219, in new_func
return func(
^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/model/generator.py", line 281, in sample_diffusion
chunk_x_l = _chunk_sample_diffusion(
^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/model/generator.py", line 225, in _chunk_sample_diffusion
x_l = tfg.step(
^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/tfg/engine.py", line 466, in step
_, grad_x0 = self._logp_and_grad_x0(
^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/tfg/engine.py", line 216, in _logp_and_grad_x0
e, g = self._energy_and_grad(x0_eps, feats, t=t, step_i=step_i)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/tfg/engine.py", line 130, in _energy_and_grad
e, g = term.energy_and_grad(coords, feats, t)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/tfg/config.py", line 273, in energy_and_grad
e, g = self._potential.energy_and_grad(coords, feats, self._params_at(t))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/tfg/potentials.py", line 129, in energy_and_grad
return self._eval(coords, feats, p, True)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/tfg/potentials.py", line 1303, in _eval
idx, eq = self._get_collision_candidates(feats)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.12/site-packages/protenix/tfg/potentials.py", line 1206, in _get_collision_candidates
return _cache_and_return(
^^^^^^^^^^^^^^^^^^
TypeError: VinaStericPotential._get_collision_candidates.._cache_and_return() takes 1 positional argument but 2 were given
```

Contributor guide

Open the contributing guide

Research direction

Reproduce the RNA-only TFG run and start in protenix/tfg/potentials.py, especially VinaStericPotential._get_collision_candidates around the failing _cache_and_return call. Trace the helper's arguments and compare the RNA-only path with the dummy-oligopeptide path. Done means RNA-only TFG inference completes without the TypeError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.