octree_attention

Bases: Octree

An enhanced Octree with transformer-specific capabilities (patching, dilation, masking).

Can be instantiated directly like a standard Octree, or created from an existing Octree instance using OctreeT.from_octree().

Once a OctreeT is instantiated, you can build the transformer context (i.e. all attention masks and relative positions) by calling the method construct_all_attention_context().

>>> octree_t = OctreeT.from_octree(octree, patch_size=26, dilation=4)  # doctest: +SKIP
>>> octree_t.construct_all_attention_context(  # doctest: +SKIP
...     nempty=True,
...     min_depth=6,
...     max_depth=10,
... )
>>> octree_t.masks[6].shape  # doctest: +SKIP
>>> octree_t.dilated_masks[6].shape  # doctest: +SKIP
>>> octree_t.rel_pos[6].shape  # doctest: +SKIP
>>> octree_t.dilated_rel_pos[6].shape  # doctest: +SKIP

Methods:

• reset –

  Resets the OctreeT to its initial state.

• from_octree –

  Creates an OctreeT from an Octree.

• construct_all_attention_context –

  Constructs all attention context for the octree.

• construct_attention_context –

  Calculates attention masks, relative positions, and padding indices

• pad_to_patch_size –

  Pads x along its first dimension so the node count at depth is a whole number of patches.

• unpad –

  Drops the padding added by pad_to_patch_size, restoring the real node count at depth.

Attributes:

• block_size (int) –

  Number of octree nodes a dilated patch spans, \(\text{patch\_size} \cdot \text{dilation}\).

block_size: int

reset() -> None

from_octree(
    octree: Octree,
    patch_size: int,
    dilation: int,
    **kwargs: Any,
) -> OctreeT

construct_all_attention_context(
    nempty: bool = False,
    min_depth: Optional[int] = None,
    max_depth: Optional[int] = None,
) -> None

construct_attention_context(
    depth: int, nempty: bool = False
) -> None

pad_to_patch_size(
    x: Tensor, depth: int, fill_value: float = 0
) -> Tensor

unpad(x: Tensor, depth: int) -> Tensor

Bases: Module

Relative Position Encoding (RPE) module used within the OctreeAttention module.

Parameters:

• patch_size (int) –

  The patch size to use for the RPE.

• num_heads (int) –

  The number of heads to use for the RPE.

• dilation (int, default: 1 ) –

  The dilation to use for the RPE.

Methods:

• pos_to_idx –

  Clamps the relative positions to \(\pm\) pos_bnd and maps each axis to its slice of the RPE table.

pos_to_idx(pos: Tensor) -> Tensor

Bases: Module

Multi-head self-attention restricted to windows of patch_size consecutive octree nodes.

A dilation above \(1\) interleaves the nodes before windowing, so a patch spans a wider region at the same cost. The attention masks and relative positions come from the OctreeT, which must have been built with construct_all_attention_context().

Parameters:

• channels (int) –

  Number of input and output channels.

• patch_size (int) –

  Number of octree nodes attending to each other.

• num_heads (int) –

  Number of attention heads.

• dilation (int, default: 1 ) –

  Stride between the nodes of a patch.

• qkv_bias (bool, default: True ) –

  Whether to use a bias in the QKV projection.

• qk_scale (Optional[float], default: None ) –

  Scaling factor for the QK matrix. Defaults to the inverse square root of the head dimension.

• attn_drop (float, default: 0.0 ) –

  Dropout rate for the attention.

• proj_drop (float, default: 0.0 ) –

  Dropout rate for the output projection.

• use_rpe (bool, default: True ) –

  Whether to add the relative position encoding to the attention logits.

Methods:

• forward_attn –

  Runs the windowed attention on patched features of shape \((N, K, C)\), returning \((N \cdot K, C)\).

• forward_proj –

  Applies the output projection and its dropout.

• forward_rpe –

  Adds the relative position bias to the attention logits, or returns them unchanged when use_rpe is off.

Attributes:

• dilated (bool) –

  Whether the patches interleave nodes rather than taking them consecutively.

dilated: bool

forward_attn(
    x: Tensor, rel_pos: Tensor, mask: Tensor
) -> Tensor

forward_proj(x: Tensor) -> Tensor

forward_rpe(
    attn: Tensor, rel_pos: Optional[Tensor] = None
) -> Tensor