initial TransformerTempFlowEstimator

pts/model/transformer_tempflow/__init__.py

@@ -0,0 +1 @@
+from .transformer_tempflow_estimator import TransformerTempFlowEstimator

pts/model/transformer_tempflow/transformer_tempflow_estimator.py

@@ -0,0 +1,234 @@
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+
+from pts import Trainer
+from pts.model import PTSEstimator, PTSPredictor, copy_parameters
+from pts.modules import RealNVP
+from pts.dataset import FieldName
+from pts.transform import (
+    Transformation,
+    Chain,
+    InstanceSplitter,
+    ExpectedNumInstanceSampler,
+    CDFtoGaussianTransform,
+    cdf_to_gaussian_forward_transform,
+    RenameFields,
+    AsNumpyArray,
+    ExpandDimArray,
+    AddObservedValuesIndicator,
+    AddTimeFeatures,
+    VstackFeatures,
+    SetFieldIfNotPresent,
+    TargetDimIndicator,
+)
+from pts.feature import (
+    TimeFeature,
+    fourier_time_features_from_frequency_str,
+    get_fourier_lags_for_frequency,
+)
+
+from .transformer_tempflow_network import TransformerTempFlowTrainingNetwork, TransformerTempFlowPredictionNetwork
+
+
+class TransformerTempFlowEstimator(PTSEstimator):
+    def __init__(
+        self,
+        input_size: int,
+        freq: str,
+        prediction_length: int,
+        target_dim: int,
+        trainer: Trainer = Trainer(),
+        context_length: Optional[int] = None,
+        d_model: int = 32,
+        dim_feedforward_scale: int = 4,
+        act_type: str = "gelu",
+        num_heads: int = 8,
+        num_encoder_layers: int = 3,
+        num_decoder_layers: int = 3,
+        num_parallel_samples: int = 100,
+        dropout_rate: float = 0.1,
+        cardinality: List[int] = [1],
+        embedding_dimension: int = 5,
+        flow_type="RealNVP",
+        n_blocks=3,
+        hidden_size=100,
+        n_hidden=2,
+        conditioning_length: int = 200,
+        dequantize: bool = False,
+
+        scaling: bool = True,
+        pick_incomplete: bool = False,
+        lags_seq: Optional[List[int]] = None,
+        time_features: Optional[List[TimeFeature]] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(trainer=trainer, **kwargs)
+
+        self.freq = freq
+        self.context_length = (
+            context_length if context_length is not None else prediction_length
+        )
+
+        self.input_size = input_size
+        self.prediction_length = prediction_length
+        self.target_dim = target_dim
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.act_type = act_type
+        self.dim_feedforward_scale = dim_feedforward_scale
+        self.num_encoder_layers = num_encoder_layers
+        self.num_decoder_layers = num_decoder_layers
+
+        self.num_parallel_samples = num_parallel_samples
+        self.dropout_rate = dropout_rate
+        self.cardinality = cardinality
+        self.embedding_dimension = embedding_dimension
+
+        self.flow_type = flow_type
+        self.n_blocks = n_blocks
+        self.hidden_size = hidden_size
+        self.n_hidden = n_hidden
+        self.conditioning_length = conditioning_length
+        self.dequantize = dequantize
+
+        self.lags_seq = (
+            lags_seq
+            if lags_seq is not None
+            else get_fourier_lags_for_frequency(freq_str=freq)
+        )
+
+        self.time_features = (
+            time_features
+            if time_features is not None
+            else fourier_time_features_from_frequency_str(self.freq)
+        )
+
+        self.history_length = self.context_length + max(self.lags_seq)
+        self.pick_incomplete = pick_incomplete
+        self.scaling = scaling
+
+    def create_transformation(self) -> Transformation:
+        return Chain(
+            [
+                AsNumpyArray(field=FieldName.TARGET, expected_ndim=2,),
+                # maps the target to (1, T)
+                # if the target data is uni dimensional
+                ExpandDimArray(field=FieldName.TARGET, axis=None,),
+                AddObservedValuesIndicator(
+                    target_field=FieldName.TARGET,
+                    output_field=FieldName.OBSERVED_VALUES,
+                ),
+                AddTimeFeatures(
+                    start_field=FieldName.START,
+                    target_field=FieldName.TARGET,
+                    output_field=FieldName.FEAT_TIME,
+                    time_features=self.time_features,
+                    pred_length=self.prediction_length,
+                ),
+                VstackFeatures(
+                    output_field=FieldName.FEAT_TIME,
+                    input_fields=[FieldName.FEAT_TIME],
+                ),
+                SetFieldIfNotPresent(field=FieldName.FEAT_STATIC_CAT, value=[0]),
+                TargetDimIndicator(
+                    field_name="target_dimension_indicator",
+                    target_field=FieldName.TARGET,
+                ),
+                AsNumpyArray(field=FieldName.FEAT_STATIC_CAT, expected_ndim=1),
+                InstanceSplitter(
+                    target_field=FieldName.TARGET,
+                    is_pad_field=FieldName.IS_PAD,
+                    start_field=FieldName.START,
+                    forecast_start_field=FieldName.FORECAST_START,
+                    train_sampler=ExpectedNumInstanceSampler(num_instances=1),
+                    past_length=self.history_length,
+                    future_length=self.prediction_length,
+                    time_series_fields=[
+                        FieldName.FEAT_TIME,
+                        FieldName.OBSERVED_VALUES,
+                    ],
+                    pick_incomplete=self.pick_incomplete,
+                ),
+                RenameFields(
+                    {
+                        f"past_{FieldName.TARGET}": f"past_{FieldName.TARGET}_cdf",
+                        f"future_{FieldName.TARGET}": f"future_{FieldName.TARGET}_cdf",
+                    }
+                ),
+            ]
+        )
+
+    def create_training_network(self, device: torch.device) -> TempFlowTrainingNetwork:
+        return TempFlowTrainingNetwork(
+            input_size=self.input_size,
+            target_dim=self.target_dim,
+            num_heads=self.num_heads,
+            act_type=self.act_type,
+            d_model=self.d_model,
+            dim_feedforward_scale=self.dim_feedforward_scale,
+            num_encoder_layers=self.num_encoder_layers,
+            num_decoder_layers=self.num_decoder_layers,
+            history_length=self.history_length,
+            context_length=self.context_length,
+            prediction_length=self.prediction_length,
+            dropout_rate=self.dropout_rate,
+            cardinality=self.cardinality,
+            embedding_dimension=self.embedding_dimension,
+            lags_seq=self.lags_seq,
+            scaling=self.scaling,
+            flow_type=self.flow_type,
+            n_blocks=self.n_blocks,
+            hidden_size=self.hidden_size,
+            n_hidden=self.n_hidden,
+            conditioning_length=self.conditioning_length,
+            dequantize=self.dequantize,
+        ).to(device)
+
+    def create_predictor(
+        self,
+        transformation: Transformation,
+        trained_network: TempFlowTrainingNetwork,
+        device: torch.device,
+    ) -> PTSPredictor:
+        prediction_network = TempFlowPredictionNetwork(
+            input_size=self.input_size,
+            target_dim=self.target_dim,
+            num_heads=self.num_heads,
+            act_type=self.act_type,
+            d_model=self.d_model,
+            dim_feedforward_scale=self.dim_feedforward_scale,
+            num_encoder_layers=self.num_encoder_layers,
+            num_decoder_layers=self.num_decoder_layers,
+            history_length=self.history_length,
+            context_length=self.context_length,
+            prediction_length=self.prediction_length,
+            dropout_rate=self.dropout_rate,
+            cardinality=self.cardinality,
+            embedding_dimension=self.embedding_dimension,
+            lags_seq=self.lags_seq,
+            scaling=self.scaling,
+            flow_type=self.flow_type,
+            n_blocks=self.n_blocks,
+            hidden_size=self.hidden_size,
+            n_hidden=self.n_hidden,
+            conditioning_length=self.conditioning_length,
+            dequantize=self.dequantize,
+            num_parallel_samples=self.num_parallel_samples,
+        ).to(device)
+
+        copy_parameters(trained_network, prediction_network)
+
+        return PTSPredictor(
+            input_transform=transformation,
+            prediction_net=prediction_network,
+            batch_size=self.trainer.batch_size,
+            freq=self.freq,
+            prediction_length=self.prediction_length,
+            device=device,
+            output_transform=None,
+        )

pts/model/transformer_tempflow/transformer_tempflow_network.py

@@ -0,0 +1,590 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+import numpy as np
+
+from pts.modules import RealNVP, MAF, FlowOutput, MeanScaler, NOPScaler
+from pts.model import weighted_average
+
+
+class TransformerTempFlowTrainingNetwork(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        d_model: int,
+        num_heads: int,
+        act_type: str,
+        dropout_rate: float,
+        dim_feedforward_scale: int,
+        num_encoder_layers: int,
+        num_decoder_layers: int,
+        history_length: int,
+        context_length: int,
+        prediction_length: int,
+        lags_seq: List[int],
+        target_dim: int,
+        conditioning_length: int,
+        flow_type: str,
+        n_blocks: int,
+        hidden_size: int,
+        n_hidden: int,
+        dequantize: bool,
+        cardinality: List[int] = [1],
+        embedding_dimension: int = 1,
+        scaling: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.target_dim = target_dim
+        self.prediction_length = prediction_length
+        self.context_length = context_length
+        self.history_length = history_length
+        self.scaling = scaling
+
+        assert len(set(lags_seq)) == len(lags_seq), "no duplicated lags allowed!"
+        lags_seq.sort()
+        self.lags_seq = lags_seq
+
+        self.encoder_input = nn.Linear(input_size, d_model)
+        self.decoder_input = nn.Linear(input_size, d_model)
+
+        # [B, T, d_model] where d_model / num_heads is int
+        self.transformer = nn.Transformer(
+            d_model=d_model,
+            nhead=num_heads,
+            num_encoder_layers=num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            dim_feedforward=dim_feedforward_scale * d_model,
+            dropout=dropout_rate,
+            activation=act_type,
+        )
+
+        flow_cls = {"RealNVP": RealNVP, "MAF": MAF,}[flow_type]
+        self.flow = flow_cls(
+            input_size=target_dim,
+            n_blocks=n_blocks,
+            n_hidden=n_hidden,
+            hidden_size=hidden_size,
+            cond_label_size=conditioning_length,
+        )
+        self.dequantize = dequantize
+
+        self.distr_output = FlowOutput(
+            self.flow, input_size=d_model, cond_size=conditioning_length
+        )
+
+        self.proj_dist_args = self.distr_output.get_args_proj(num_cells)
+
+        self.embed_dim = 1
+        self.embed = nn.Embedding(
+            num_embeddings=self.target_dim, embedding_dim=self.embed_dim
+        )
+
+        self.scaler = MeanScaler(keepdim=True)
+
+    @staticmethod
+    def get_lagged_subsequences(
+        sequence: torch.Tensor,
+        sequence_length: int,
+        indices: List[int],
+        subsequences_length: int = 1,
+    ) -> torch.Tensor:
+        """
+        Returns lagged subsequences of a given sequence.
+        Parameters
+        ----------
+        sequence
+            the sequence from which lagged subsequences should be extracted.
+            Shape: (N, T, C).
+        sequence_length
+            length of sequence in the T (time) dimension (axis = 1).
+        indices
+            list of lag indices to be used.
+        subsequences_length
+            length of the subsequences to be extracted.
+        Returns
+        --------
+        lagged : Tensor
+            a tensor of shape (N, S, C, I),
+            where S = subsequences_length and I = len(indices),
+            containing lagged subsequences.
+            Specifically, lagged[i, :, j, k] = sequence[i, -indices[k]-S+j, :].
+        """
+        # we must have: history_length + begin_index >= 0
+        # that is: history_length - lag_index - sequence_length >= 0
+        # hence the following assert
+        assert max(indices) + subsequences_length <= sequence_length, (
+            f"lags cannot go further than history length, found lag "
+            f"{max(indices)} while history length is only {sequence_length}"
+        )
+        assert all(lag_index >= 0 for lag_index in indices)
+
+        lagged_values = []
+        for lag_index in indices:
+            begin_index = -lag_index - subsequences_length
+            end_index = -lag_index if lag_index > 0 else None
+            lagged_values.append(sequence[:, begin_index:end_index, ...].unsqueeze(1))
+        return torch.cat(lagged_values, dim=1).permute(0, 2, 3, 1)
+
+    def unroll(
+        self,
+        lags: torch.Tensor,
+        scale: torch.Tensor,
+        time_feat: torch.Tensor,
+        target_dimension_indicator: torch.Tensor,
+        unroll_length: int,
+        begin_state: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+    ) -> Tuple[
+        torch.Tensor,
+        Union[List[torch.Tensor], torch.Tensor],
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+
+        # (batch_size, sub_seq_len, target_dim, num_lags)
+        lags_scaled = lags / scale.unsqueeze(-1)
+
+        # assert_shape(
+        #     lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)),
+        # )
+
+        input_lags = lags_scaled.reshape(
+            (-1, unroll_length, len(self.lags_seq) * self.target_dim)
+        )
+
+        # (batch_size, target_dim, embed_dim)
+        index_embeddings = self.embed(target_dimension_indicator)
+        # assert_shape(index_embeddings, (-1, self.target_dim, self.embed_dim))
+
+        # (batch_size, seq_len, target_dim * embed_dim)
+        repeated_index_embeddings = (
+            index_embeddings.unsqueeze(1)
+            .expand(-1, unroll_length, -1, -1)
+            .reshape((-1, unroll_length, self.target_dim * self.embed_dim))
+        )
+
+        # (batch_size, sub_seq_len, input_dim)
+        inputs = torch.cat((input_lags, repeated_index_embeddings, time_feat), dim=-1)
+
+        # unroll encoder
+        outputs, state = self.rnn(inputs, begin_state)
+
+        # assert_shape(outputs, (-1, unroll_length, self.num_cells))
+        # for s in state:
+        #     assert_shape(s, (-1, self.num_cells))
+
+        # assert_shape(
+        #     lags_scaled, (-1, unroll_length, self.target_dim, len(self.lags_seq)),
+        # )
+
+        return outputs, state, lags_scaled, inputs
+
+    def unroll_encoder(
+        self,
+        past_time_feat: torch.Tensor,
+        past_target_cdf: torch.Tensor,
+        past_observed_values: torch.Tensor,
+        past_is_pad: torch.Tensor,
+        future_time_feat: Optional[torch.Tensor],
+        future_target_cdf: Optional[torch.Tensor],
+        target_dimension_indicator: torch.Tensor,
+    ) -> Tuple[
+        torch.Tensor,
+        Union[List[torch.Tensor], torch.Tensor],
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+    ]:
+        """
+        Unrolls the RNN encoder over past and, if present, future data.
+        Returns outputs and state of the encoder, plus the scale of
+        past_target_cdf and a vector of static features that was constructed
+        and fed as input to the encoder. All tensor arguments should have NTC
+        layout.
+
+        Parameters
+        ----------
+        past_time_feat
+            Past time features (batch_size, history_length, num_features)
+        past_target_cdf
+            Past marginal CDF transformed target values (batch_size,
+            history_length, target_dim)
+        past_observed_values
+            Indicator whether or not the values were observed (batch_size,
+            history_length, target_dim)
+        past_is_pad
+            Indicator whether the past target values have been padded
+            (batch_size, history_length)
+        future_time_feat
+            Future time features (batch_size, prediction_length, num_features)
+        future_target_cdf
+            Future marginal CDF transformed target values (batch_size,
+            prediction_length, target_dim)
+        target_dimension_indicator
+            Dimensionality of the time series (batch_size, target_dim)
+
+        Returns
+        -------
+        outputs
+            RNN outputs (batch_size, seq_len, num_cells)
+        states
+            RNN states. Nested list with (batch_size, num_cells) tensors with
+        dimensions target_dim x num_layers x (batch_size, num_cells)
+        scale
+            Mean scales for the time series (batch_size, 1, target_dim)
+        lags_scaled
+            Scaled lags(batch_size, sub_seq_len, target_dim, num_lags)
+        inputs
+            inputs to the RNN
+
+        """
+
+        past_observed_values = torch.min(
+            past_observed_values, 1 - past_is_pad.unsqueeze(-1)
+        )
+
+        if future_time_feat is None or future_target_cdf is None:
+            time_feat = past_time_feat[:, -self.context_length :, ...]
+            sequence = past_target_cdf
+            sequence_length = self.history_length
+            subsequences_length = self.context_length
+        else:
+            time_feat = torch.cat(
+                (past_time_feat[:, -self.context_length :, ...], future_time_feat),
+                dim=1,
+            )
+            sequence = torch.cat((past_target_cdf, future_target_cdf), dim=1)
+            sequence_length = self.history_length + self.prediction_length
+            subsequences_length = self.context_length + self.prediction_length
+
+        # (batch_size, sub_seq_len, target_dim, num_lags)
+        lags = self.get_lagged_subsequences(
+            sequence=sequence,
+            sequence_length=sequence_length,
+            indices=self.lags_seq,
+            subsequences_length=subsequences_length,
+        )
+
+        # scale is computed on the context length last units of the past target
+        # scale shape is (batch_size, 1, target_dim)
+        _, scale = self.scaler(
+            past_target_cdf[:, -self.context_length :, ...],
+            past_observed_values[:, -self.context_length :, ...],
+        )
+
+        outputs, states, lags_scaled, inputs = self.unroll(
+            lags=lags,
+            scale=scale,
+            time_feat=time_feat,
+            target_dimension_indicator=target_dimension_indicator,
+            unroll_length=subsequences_length,
+            begin_state=None,
+        )
+
+        return outputs, states, scale, lags_scaled, inputs
+
+    def distr_args(self, rnn_outputs: torch.Tensor):
+        """
+        Returns the distribution of DeepVAR with respect to the RNN outputs.
+
+        Parameters
+        ----------
+        rnn_outputs
+            Outputs of the unrolled RNN (batch_size, seq_len, num_cells)
+        scale
+            Mean scale for each time series (batch_size, 1, target_dim)
+
+        Returns
+        -------
+        distr
+            Distribution instance
+        distr_args
+            Distribution arguments
+        """
+        (distr_args,) = self.proj_dist_args(rnn_outputs)
+
+        # # compute likelihood of target given the predicted parameters
+        # distr = self.distr_output.distribution(distr_args, scale=scale)
+
+        # return distr, distr_args
+        return distr_args
+
+    def forward(
+        self,
+        target_dimension_indicator: torch.Tensor,
+        past_time_feat: torch.Tensor,
+        past_target_cdf: torch.Tensor,
+        past_observed_values: torch.Tensor,
+        past_is_pad: torch.Tensor,
+        future_time_feat: torch.Tensor,
+        future_target_cdf: torch.Tensor,
+        future_observed_values: torch.Tensor,
+    ) -> Tuple[torch.Tensor, ...]:
+        """
+        Computes the loss for training DeepVAR, all inputs tensors representing
+        time series have NTC layout.
+
+        Parameters
+        ----------
+        target_dimension_indicator
+            Indices of the target dimension (batch_size, target_dim)
+        past_time_feat
+            Dynamic features of past time series (batch_size, history_length,
+            num_features)
+        past_target_cdf
+            Past marginal CDF transformed target values (batch_size,
+            history_length, target_dim)
+        past_observed_values
+            Indicator whether or not the values were observed (batch_size,
+            history_length, target_dim)
+        past_is_pad
+            Indicator whether the past target values have been padded
+            (batch_size, history_length)
+        future_time_feat
+            Future time features (batch_size, prediction_length, num_features)
+        future_target_cdf
+            Future marginal CDF transformed target values (batch_size,
+            prediction_length, target_dim)
+        future_observed_values
+            Indicator whether or not the future values were observed
+            (batch_size, prediction_length, target_dim)
+
+        Returns
+        -------
+        distr
+            Loss with shape (batch_size, 1)
+        likelihoods
+            Likelihoods for each time step
+            (batch_size, context + prediction_length, 1)
+        distr_args
+            Distribution arguments (context + prediction_length,
+            number_of_arguments)
+        """
+
+        seq_len = self.context_length + self.prediction_length
+
+        # unroll the decoder in "training mode", i.e. by providing future data
+        # as well
+        rnn_outputs, _, scale, _, _ = self.unroll_encoder(
+            past_time_feat=past_time_feat,
+            past_target_cdf=past_target_cdf,
+            past_observed_values=past_observed_values,
+            past_is_pad=past_is_pad,
+            future_time_feat=future_time_feat,
+            future_target_cdf=future_target_cdf,
+            target_dimension_indicator=target_dimension_indicator,
+        )
+
+        # put together target sequence
+        # (batch_size, seq_len, target_dim)
+        target = torch.cat(
+            (past_target_cdf[:, -self.context_length :, ...], future_target_cdf), dim=1,
+        )
+
+        # assert_shape(target, (-1, seq_len, self.target_dim))
+
+        distr_args = self.distr_args(rnn_outputs=rnn_outputs)
+        if self.scaling:
+            self.flow.scale = scale
+
+        # we sum the last axis to have the same shape for all likelihoods
+        # (batch_size, subseq_length, 1)
+        if self.dequantize:
+            target += torch.rand_like(target)
+        likelihoods = -self.flow.log_prob(target, distr_args).unsqueeze(-1)
+
+        # assert_shape(likelihoods, (-1, seq_len, 1))
+
+        past_observed_values = torch.min(
+            past_observed_values, 1 - past_is_pad.unsqueeze(-1)
+        )
+
+        # (batch_size, subseq_length, target_dim)
+        observed_values = torch.cat(
+            (
+                past_observed_values[:, -self.context_length :, ...],
+                future_observed_values,
+            ),
+            dim=1,
+        )
+
+        # mask the loss at one time step if one or more observations is missing
+        # in the target dimensions (batch_size, subseq_length, 1)
+        loss_weights, _ = observed_values.min(dim=-1, keepdim=True)
+
+        # assert_shape(loss_weights, (-1, seq_len, 1))
+
+        loss = weighted_average(likelihoods, weights=loss_weights, dim=1)
+
+        # assert_shape(loss, (-1, -1, 1))
+
+        # self.distribution = distr
+
+        return (loss.mean(), likelihoods, distr_args)
+
+
+class TransformerTempFlowPredictionNetwork(TransformerTempFlowTrainingNetwork):
+    def __init__(self, num_parallel_samples: int, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.num_parallel_samples = num_parallel_samples
+
+        # for decoding the lags are shifted by one,
+        # at the first time-step of the decoder a lag of one corresponds to
+        # the last target value
+        self.shifted_lags = [l - 1 for l in self.lags_seq]
+
+    def sampling_decoder(
+        self,
+        past_target_cdf: torch.Tensor,
+        target_dimension_indicator: torch.Tensor,
+        time_feat: torch.Tensor,
+        scale: torch.Tensor,
+        begin_states: Union[List[torch.Tensor], torch.Tensor],
+    ) -> torch.Tensor:
+        """
+        Computes sample paths by unrolling the RNN starting with a initial
+        input and state.
+
+        Parameters
+        ----------
+        past_target_cdf
+            Past marginal CDF transformed target values (batch_size,
+            history_length, target_dim)
+        target_dimension_indicator
+            Indices of the target dimension (batch_size, target_dim)
+        time_feat
+            Dynamic features of future time series (batch_size, history_length,
+            num_features)
+        scale
+            Mean scale for each time series (batch_size, 1, target_dim)
+        begin_states
+            List of initial states for the RNN layers (batch_size, num_cells)
+        Returns
+        --------
+        sample_paths : Tensor
+            A tensor containing sampled paths. Shape: (1, num_sample_paths,
+            prediction_length, target_dim).
+        """
+
+        def repeat(tensor, dim=0):
+            return tensor.repeat_interleave(repeats=self.num_parallel_samples, dim=dim)
+
+        # blows-up the dimension of each tensor to
+        # batch_size * self.num_sample_paths for increasing parallelism
+        repeated_past_target_cdf = repeat(past_target_cdf)
+        repeated_time_feat = repeat(time_feat)
+        repeated_scale = repeat(scale)
+        if self.scaling:
+            self.flow.scale = repeated_scale
+        repeated_target_dimension_indicator = repeat(target_dimension_indicator)
+
+        if self.cell_type == "LSTM":
+            repeated_states = [repeat(s, dim=1) for s in begin_states]
+        else:
+            repeated_states = repeat(begin_states, dim=1)
+
+        future_samples = []
+
+        # for each future time-units we draw new samples for this time-unit
+        # and update the state
+        for k in range(self.prediction_length):
+            lags = self.get_lagged_subsequences(
+                sequence=repeated_past_target_cdf,
+                sequence_length=self.history_length + k,
+                indices=self.shifted_lags,
+                subsequences_length=1,
+            )
+
+            rnn_outputs, repeated_states, _, _ = self.unroll(
+                begin_state=repeated_states,
+                lags=lags,
+                scale=repeated_scale,
+                time_feat=repeated_time_feat[:, k : k + 1, ...],
+                target_dimension_indicator=repeated_target_dimension_indicator,
+                unroll_length=1,
+            )
+
+            distr_args = self.distr_args(rnn_outputs=rnn_outputs)
+
+            # (batch_size, 1, target_dim)
+            new_samples = self.flow.sample(cond=distr_args)
+
+            # (batch_size, seq_len, target_dim)
+            future_samples.append(new_samples)
+            repeated_past_target_cdf = torch.cat(
+                (repeated_past_target_cdf, new_samples), dim=1
+            )
+
+        # (batch_size * num_samples, prediction_length, target_dim)
+        samples = torch.cat(future_samples, dim=1)
+
+        # (batch_size, num_samples, prediction_length, target_dim)
+        return samples.reshape(
+            (-1, self.num_parallel_samples, self.prediction_length, self.target_dim,)
+        )
+
+    def forward(
+        self,
+        target_dimension_indicator: torch.Tensor,
+        past_time_feat: torch.Tensor,
+        past_target_cdf: torch.Tensor,
+        past_observed_values: torch.Tensor,
+        past_is_pad: torch.Tensor,
+        future_time_feat: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Predicts samples given the trained DeepVAR model.
+        All tensors should have NTC layout.
+        Parameters
+        ----------
+        target_dimension_indicator
+            Indices of the target dimension (batch_size, target_dim)
+        past_time_feat
+            Dynamic features of past time series (batch_size, history_length,
+            num_features)
+        past_target_cdf
+            Past marginal CDF transformed target values (batch_size,
+            history_length, target_dim)
+        past_observed_values
+            Indicator whether or not the values were observed (batch_size,
+            history_length, target_dim)
+        past_is_pad
+            Indicator whether the past target values have been padded
+            (batch_size, history_length)
+        future_time_feat
+            Future time features (batch_size, prediction_length, num_features)
+
+        Returns
+        -------
+        sample_paths : Tensor
+            A tensor containing sampled paths (1, num_sample_paths,
+            prediction_length, target_dim).
+
+        """
+
+        # mark padded data as unobserved
+        # (batch_size, target_dim, seq_len)
+        past_observed_values = torch.min(
+            past_observed_values, 1 - past_is_pad.unsqueeze(-1)
+        )
+
+        # unroll the decoder in "prediction mode", i.e. with past data only
+        _, begin_states, scale, _, _ = self.unroll_encoder(
+            past_time_feat=past_time_feat,
+            past_target_cdf=past_target_cdf,
+            past_observed_values=past_observed_values,
+            past_is_pad=past_is_pad,
+            future_time_feat=None,
+            future_target_cdf=None,
+            target_dimension_indicator=target_dimension_indicator,
+        )
+
+        return self.sampling_decoder(
+            past_target_cdf=past_target_cdf,
+            target_dimension_indicator=target_dimension_indicator,
+            time_feat=future_time_feat,
+            scale=scale,
+            begin_states=begin_states,
+        )