Skip to content

boxmot.segmentors.exporters.edgetam.attention

Exportable attention with real rotary arithmetic and explicit memory masks.

ExportAttention 

ExportAttention(source: Any)

Bases: Module

Reuse trained SAM projections without CUDA context managers during export.

k_proj instance-attribute 

k_proj = k_proj

num_heads instance-attribute 

num_heads = num_heads

out_proj instance-attribute 

out_proj = out_proj

q_proj instance-attribute 

q_proj = q_proj

v_proj instance-attribute 

v_proj = v_proj

combine 

combine(x: Tensor) -> Tensor

Restore projected channels after attention.

forward 

forward(q: Tensor, k: Tensor, v: Tensor) -> Tensor

Apply scaled dot product attention with unchanged learned projections.

heads 

heads(x: Tensor) -> Tensor

Convert batch/token/channel tensors into attention heads.

RotaryAttention 

RotaryAttention(
    source: Any,
    *,
    image_size: int,
    memory_slots: int = 0,
    tokens_per_memory: int = 0,
)

Bases: ExportAttention

Real-valued RoPE attention for fixed image and padded memory token grids.

memory_slots instance-attribute 

memory_slots = memory_slots

tokens_per_memory instance-attribute 

tokens_per_memory = tokens_per_memory

forward 

forward(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    valid: Tensor | None = None,
) -> Tensor

Exclude absent slots and leave global memories and object pointers unrotated.

MemoryAttentionLayer 

MemoryAttentionLayer(
    source: Any,
    *,
    image_size: int,
    memory_slots: int,
    tokens_per_memory: int,
)

Bases: Module

Preserve an official memory layer while masking absent history entries.

cross_attention instance-attribute 

cross_attention = RotaryAttention(
    cross_attn_image,
    image_size=image_size,
    memory_slots=memory_slots,
    tokens_per_memory=tokens_per_memory,
)

self_attention instance-attribute 

self_attention = RotaryAttention(
    self_attn, image_size=image_size
)

source instance-attribute 

source = source

forward 

forward(
    x: Tensor,
    memory: Tensor,
    image_pos: Tensor,
    memory_pos: Tensor,
    valid: Tensor,
) -> Tensor

Attend to valid histories without changing their softmax normalization.

MemoryAttention 

MemoryAttention(
    source: Any,
    *,
    image_size: int,
    memory_slots: int,
    tokens_per_memory: int,
)

Bases: Module

Batch-first memory attention accepting masked, fixed-capacity histories.

layers instance-attribute 

layers = ModuleList(
    (
        MemoryAttentionLayer(
            layer,
            image_size=image_size,
            memory_slots=memory_slots,
            tokens_per_memory=tokens_per_memory,
        )
    )
    for layer in (layers)
)

norm instance-attribute 

norm = norm

pos_enc_at_input instance-attribute 

pos_enc_at_input = pos_enc_at_input

forward 

forward(
    image: Tensor,
    memory: Tensor,
    image_pos: Tensor,
    memory_pos: Tensor,
    valid: Tensor,
) -> Tensor

Fuse one independently masked history with each object's current features.

attend 

attend(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    valid: Tensor | None = None,
) -> Tensor

Evaluate attention in query blocks to bound the eager attention workspace.

real_rotary 

real_rotary(
    x: Tensor, cosine: Tensor, sine: Tensor
) -> Tensor

Apply the upstream interleaved complex rotation using real operations.