Implementation:Sktime Pytorch forecasting FullAttention

Knowledge Sources	Sktime_Pytorch_forecasting
Domains	Time_Series, Forecasting, Deep_Learning
Last Updated	2026-02-08 08:00 GMT

Overview

Scaled dot-product attention implementation with optional causal masking, dropout, and an efficient PyTorch SDPA backend.

Description

The FullAttention class implements the standard scaled dot-product attention mechanism as an nn.Module. It supports two computation paths: a traditional einsum-based path that manually computes attention scores, applies optional triangular causal masking, and performs softmax normalization with dropout; and an efficient attention path that delegates to PyTorch's native scaled_dot_product_attention function, which automatically selects the optimal backend (FlashAttention-2, Memory-Efficient Attention, or C++ implementation) based on hardware and input properties.

The module also provides TriangularCausalMask, a helper class that generates an upper-triangular boolean mask of shape (B, 1, L, L) to prevent attending to future positions in autoregressive decoding scenarios. The output_attention flag cannot be combined with use_efficient_attention since the efficient backend does not return attention weights.

Usage

Use FullAttention as the inner attention mechanism for AttentionLayer in transformer-based forecasting models. Enable use_efficient_attention=True for faster inference on long sequences, or use the default einsum path when attention weight visualization is needed.

Code Reference

Source Location

• Repository: Sktime_Pytorch_forecasting
• File: pytorch_forecasting/layers/_attention/_full_attention.py
• Lines: 1-117

Signatures

TriangularCausalMask

class TriangularCausalMask:
    def __init__(self, B, L, device="cpu")
    @property
    def mask(self)

FullAttention

class FullAttention(nn.Module):
    def __init__(
        self,
        mask_flag=True,
        factor=5,
        scale=None,
        attention_dropout=0.1,
        output_attention=False,
        use_efficient_attention=False,
    )
    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None)

Import

from pytorch_forecasting.layers._attention._full_attention import FullAttention, TriangularCausalMask

I/O Contract

Inputs

Name	Type	Required	Description
mask_flag	bool	No	Whether to apply causal masking (default: True)
factor	int	No	Factor for scaling attention scores (default: 5)
scale	float or None	No	Custom scaling factor for attention scores. None uses 1/sqrt(E)
attention_dropout	float	No	Dropout rate applied to attention scores (default: 0.1)
output_attention	bool	No	Whether to return attention weights (default: False). Cannot be True with use_efficient_attention
use_efficient_attention	bool	No	Whether to use PyTorch native SDPA backend for faster computation (default: False)
queries	torch.Tensor	Yes	Query tensor of shape (B, L, H, E)
keys	torch.Tensor	Yes	Key tensor of shape (B, S, H, E)
values	torch.Tensor	Yes	Value tensor of shape (B, S, H, D)
attn_mask	TriangularCausalMask or None	Yes	Attention mask object, or None for automatic causal mask generation
B	int	Yes	Batch size for TriangularCausalMask
L	int	Yes	Sequence length for TriangularCausalMask
device	str	No	Device for mask tensor (default: "cpu")

Outputs

Name	Type	Description
V	torch.Tensor	Attended values tensor of shape (B, L, H, D), returned as contiguous
A	torch.Tensor or None	Attention weights of shape (B, H, L, S) if output_attention=True, otherwise None
mask	torch.BoolTensor	Upper-triangular boolean mask of shape (B, 1, L, L) from TriangularCausalMask.mask property

Usage Examples

import torch
from pytorch_forecasting.layers._attention._full_attention import FullAttention, TriangularCausalMask

# Standard attention with causal masking
attention = FullAttention(
    mask_flag=True,
    attention_dropout=0.1,
    output_attention=False,
)

B, L, S, H, E, D = 32, 96, 96, 8, 64, 64
queries = torch.randn(B, L, H, E)
keys = torch.randn(B, S, H, E)
values = torch.randn(B, S, H, D)

output, attn = attention(queries, keys, values, attn_mask=None)
# output shape: (32, 96, 8, 64)

# Efficient attention using PyTorch native SDPA
efficient_attention = FullAttention(
    mask_flag=True,
    use_efficient_attention=True,
    attention_dropout=0.1,
)
output, _ = efficient_attention(queries, keys, values, attn_mask=None)

# Create a causal mask manually
causal_mask = TriangularCausalMask(B=32, L=96, device="cpu")
print(causal_mask.mask.shape)  # torch.Size([32, 1, 96, 96])

Related Pages

• Principle:Sktime_Pytorch_forecasting_Scaled_Dot_Product_Attention