deepfold.ops package

Submodules

deepfold.ops.evoformer_attention module

deepfold.ops.evoformer_attention.deepspeed_evo_attn(q: Tensor, k: Tensor, v: Tensor, biases: List[Tensor]) → Tensor

“” Compute attention using the DeepSpeed DS4Sci_EvoformerAttention kernel.

Parameters:

• q – [*, H, Q, C_hidden] query data

• k – [*, H, K, C_hidden] key data

• v – [*, H, V, C_hidden] value data

• biases – List of biases that broadcast to [*, H, Q, K]

Module contents