Implementation:Sktime Pytorch forecasting FlattenHead

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

Overview

FlattenHead is an output head layer that flattens hidden states and projects them to the prediction window via a linear transformation.

Description

FlattenHead extends nn.Module and serves as the final output layer for models such as TimeXer. It takes a multi-dimensional hidden representation, flattens the last two dimensions, applies a linear projection to the target window size, and optionally reshapes the output for quantile predictions. Dropout regularization is applied after the linear layer.

Usage

Use FlattenHead when building a forecasting model that needs to convert encoder hidden states into a flat prediction tensor of shape matching the forecast horizon. It supports both point forecasts and quantile forecasts via the optional n_quantiles parameter.

Code Reference

Source Location

• Repository: Sktime_Pytorch_forecasting
• File: pytorch_forecasting/layers/_output/_flatten_head.py
• Lines: 1-46

Signature

class FlattenHead(nn.Module):
    def __init__(self, n_vars, nf, target_window, head_dropout=0, n_quantiles=None):
    def forward(self, x):

Import

from pytorch_forecasting.layers._output._flatten_head import FlattenHead

I/O Contract

Inputs

__init__

Name	Type	Required	Description
n_vars	int	Yes	Number of input features (variables).
nf	int	Yes	Number of features in the last hidden layer.
target_window	int	Yes	Target prediction window size (forecast horizon).
head_dropout	float	No	Dropout rate applied after the linear layer. Defaults to 0.
n_quantiles	int or None	No	Number of quantiles for quantile regression. Defaults to None (point forecast).

forward

Name	Type	Required	Description
x	torch.Tensor	Yes	Hidden state tensor to be flattened and projected.

Outputs

forward

Name	Type	Description
x	torch.Tensor	Prediction tensor. Shape is (batch_size, target_window, n_vars) for point forecasts, or (batch_size, target_window, n_quantiles) for quantile forecasts.

Usage Examples

import torch
from pytorch_forecasting.layers._output._flatten_head import FlattenHead

# Point forecast head
head = FlattenHead(n_vars=7, nf=512, target_window=24, head_dropout=0.1)
hidden = torch.randn(32, 7, 512)  # (batch, n_vars, nf)
output = head(hidden)
# output shape: (32, 24, 7)

# Quantile forecast head
head_q = FlattenHead(n_vars=7, nf=512, target_window=24, head_dropout=0.1, n_quantiles=3)
output_q = head_q(hidden)
# output_q shape: (32, 24, 3)

Related Pages

• Principle:Sktime_Pytorch_forecasting_TimeXer_Architecture