TODO list

[OccupyMars2030]

# protenix\model\protenix.py  line 337-349
# def main_inference_loop(

        for _ in range(N_model_seed):
            pred_dict, log_dict, time_tracker = self._main_inference_loop(
                input_feature_dict=input_feature_dict,
                label_dict=label_dict,
                N_cycle=N_cycle,
                mode=mode,
                inplace_safe=inplace_safe,
                chunk_size=chunk_size,
                symmetric_permutation=symmetric_permutation,
            )
            pred_dicts.append(pred_dict)
            log_dicts.append(log_dict)
            time_trackers.append(time_tracker)

In each iteration of for loop, we should use a different seed, but where are the different seeds ?

[OccupyMars2030]

2025/9/7 23:32

!protenix predict --input examples/example.json --out_dir  ./output --seeds 101 --model_name "protenix_base_default_v0.5.0"


# protenix\model\protenix.py  line:372
    # def _main_inference_loop(
    #     self,
    #     input_feature_dict: dict[str, Any],
    #     label_dict: dict[str, Any],
    #     N_cycle: int,
    #     mode: str,
    #     inplace_safe: bool = True,
    #     chunk_size: Optional[int] = 4,
    #     symmetric_permutation: SymmetricPermutation = None,
    # ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, Any]]:
    #     """
    #     Main inference loop (single model seed) for the Alphafold3 model.

    #     Returns:
    #         tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, Any]]: Prediction, log, and time dictionaries.
    #     """

    #     inspect_input_features(input_feature_dict)
    #     exit()

I write a function inspect_input_features to check input_feature_dict just as the commented code shows. Then run the above command which is copied from Protenix README.md.

---

bytedance#188
bytedance#189
bytedance#190

=== bond_mask ===
Shape: torch.Size([2529, 2529])
Dtype: torch.int64
Device: cuda:0
Min: 0
Max: 0

bond_mask in input_feature_dict is a huge tensor which occupies too much memory (its size is [Natom, Natom]), what is its function ? Can we reduce its size ?
Note: bond_mask is not defined in the AF3 paper.

[OccupyMars2030]

study why it is needed to implement LayerNorm with custom CUDA code

protenix\openfold_local\model\primitives.py, on branch my-v0.1.0-study-notes

# Keep the function name for code simplicity
def LayerNorm(c_in, eps: float = 1e-5):
    # if specify "fast_layernorm" and fastln_is_installed, use the FusedLayerNorm,
    # Otherwise, OpenFoldLayerNorm is used!
    if fastln_is_installed:
        # print("use fast layernorm")
        return FusedLayerNorm(c_in, eps)
    # print("use openfold layernorm")
    return OpenFoldLayerNorm(c_in, eps)