openfold.model.triangular_attention

Classes

TriangleAttention(c_in, c_hidden, no_heads)
TriangleAttentionEndingNode(c_in, c_hidden, ...)	Implements Algorithm 14.
TriangleAttentionStartingNode	alias of TriangleAttention

class TriangleAttention(c_in, c_hidden, no_heads, starting=True, inf=1000000000.0)

Bases: Module

forward(x, mask=None, chunk_size=None, use_memory_efficient_kernel=False, use_deepspeed_evo_attention=False, use_cuequivariance_attention=False, use_lma=False, inplace_safe=False)

Parameters:

• x (Tensor) – [*, I, J, C_in] input tensor (e.g. the pair representation)

• mask (Tensor | None)

• chunk_size (int | None)

• use_memory_efficient_kernel (bool)

• use_deepspeed_evo_attention (bool)

• use_cuequivariance_attention (bool)

• use_lma (bool)

• inplace_safe (bool)

Returns:

[*, I, J, C_in] output tensor

Return type:

Tensor

class TriangleAttentionEndingNode(c_in, c_hidden, no_heads, *, starting=False, inf=1000000000.0)

Bases: TriangleAttention

Implements Algorithm 14.

TriangleAttentionStartingNode

alias of TriangleAttention