content: __all__ = ["DistributedConfig", "BaseModel"]
  
  
  import inspect
  import math
  import random
  import warnings
  from contextlib import contextmanager
  from copy import deepcopy
  from dataclasses import dataclass
  from typing import Dict, List, Union
  
  import fsspec
  import numpy as np
  import pytorch_lightning as pl
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from pytorch_lightning.callbacks.early_stopping import EarlyStopping
  
  import neuralforecast.losses.pytorch as losses
  from neuralforecast.common.enums import ExplainerEnum
  from neuralforecast.tsdataset import (
      BaseTimeSeriesDataset,
      TimeSeriesDataModule,
      _DistributedTimeSeriesDataModule,
  )
  
  from ..losses.pytorch import BasePointLoss, DistributionLoss
  from ..utils import get_indexer_raise_missing
  
  DISTRIBUTION_LOSSES = (
      losses.DistributionLoss,
      losses.PMM,
      losses.GMM,
      losses.NBMM,
  )
  MULTIQUANTILE_LOSSES = (
      losses.MQLoss,
      losses.HuberMQLoss,
  )
  from ._scalers import TemporalNorm
  
  
  @dataclass
  class DistributedConfig:
      partitions_path: str
      num_nodes: int
      devices: int
  
  
  @contextmanager
  def _disable_torch_init():
      """Context manager used to disable pytorch's weight initialization.
  
      This is especially useful when loading saved models, since when initializing
      a model the weights are also initialized following some method
      (e.g. kaiming uniform), and that time is wasted since we'll override them with
      the saved weights."""
  
      def noop(*args, **kwargs):
          return
  
      kaiming_uniform = nn.init.kaiming_uniform_
      kaiming_normal = nn.init.kaiming_normal_
      xavier_uniform = nn.init.xavier_uniform_
      xavier_normal = nn.init.xavier_normal_
  
      nn.init.kaiming_uniform_ = noop
      nn.init.kaiming_normal_ = noop
      nn.init.xavier_uniform_ = noop
      nn.init.xavier_normal_ = noop
      try:
          yield
      finally:
          nn.init.kaiming_uniform_ = kaiming_uniform
          nn.init.kaiming_normal_ = kaiming_normal
          nn.init.xavier_uniform_ = xavier_uniform
          nn.init.xavier_normal_ = xavier_normal
  
  
  def tensor_to_numpy(tensor: torch.Tensor) -> np.ndarray:
      """Convert a tensor to numpy"""
      if tensor.dtype == torch.bfloat16:
          return tensor.float().numpy()
  
      return tensor.numpy()
  
  
  class BaseModel(pl.LightningModule):
      EXOGENOUS_FUTR = True  # If the model can handle future exogenous variables
      EXOGENOUS_HIST = True  # If the model can handle historical exogenous variables
      EXOGENOUS_STAT = True  # If the model can handle static exogenous variables
      MULTIVARIATE = False  # If the model produces multivariate forecasts (True) or univariate (False)
      RECURRENT = (
          False  # If the model produces forecasts recursively (True) or direct (False)
      )
  
      def __init__(
          self,
          h: int,
          input_size: int,
          loss: Union[BasePointLoss, DistributionLoss, nn.Module],
          valid_loss: Union[BasePointLoss, DistributionLoss, nn.Module],
          learning_rate: float,
          max_steps: int,
          val_check_steps: int,
          batch_size: int,
          valid_batch_size: Union[int, None],
          windows_batch_size: int,
          inference_windows_batch_size: Union[int, None],
          start_padding_enabled: bool,
          training_data_availability_threshold: Union[float, List[float]] = 0.0,
          n_series: Union[int, None] = None,
          n_samples: Union[int, None] = 100,
          h_train: int = 1,
          inference_input_size: Union[int, None] = None,
          step_size: int = 1,
          num_lr_decays: int = 0,
          early_stop_patience_steps: int = -1,
          scaler_type: str = "identity",
          futr_exog_list: Union[List, None] = None,
          hist_exog_list: Union[List, None] = None,
          stat_exog_list: Union[List, None] = None,
          exclude_insample_y: Union[bool, None] = False,
          drop_last_loader: Union[bool, None] = False,
          random_seed: Union[int, None] = 1,
          alias: Union[str, None] = None,
          optimizer: Union[torch.optim.Optimizer, None] = None,
          optimizer_kwargs: Union[Dict, None] = None,
          lr_scheduler: Union[torch.optim.lr_scheduler.LRScheduler, None] = None,
          lr_scheduler_kwargs: Union[Dict, None] = None,
          dataloader_kwargs=None,
          **trainer_kwargs,
      ):
          super().__init__()
  
          # Multivarariate checks
          if self.MULTIVARIATE and n_series is None:
              raise Exception(
                  f"{type(self).__name__} is a multivariate model. Please set n_series to the number of unique time series in your dataset."
              )
          if not self.MULTIVARIATE:
              n_series = 1
          self.n_series = n_series
          self.n_predicts = 1
  
          # Protections for previous recurrent models
          if input_size < 1:
              input_size = 3 * h
              warnings.warn(
                  f"Input size too small. Automatically setting input size to 3 * horizon = {input_size}"
              )
  
          if inference_input_size is None:
              inference_input_size = input_size
          elif inference_input_size is not None and inference_input_size < 1:
              inference_input_size = input_size
              warnings.warn(
                  f"Inference input size too small. Automatically setting inference input size to input_size = {input_size}"
              )
  
          # For recurrent models we need one additional input as we need to shift insample_y to use it as input
          if self.RECURRENT:
              input_size += 1
              inference_input_size += 1
  
          # Attributes needed for recurrent models
          self.horizon_backup = h
          self.input_size_backup = input_size
          self.n_samples = n_samples
          if self.RECURRENT:
              if (
                  hasattr(loss, "horizon_weight")
                  and loss.horizon_weight is not None
                  and h_train != h
              ):
                  warnings.warn(
                      f"Setting h_train={h} to match the horizon_weight length."
                  )
                  h_train = h
              self.h_train = h_train
              self.inference_input_size = inference_input_size
              self.rnn_state = None
              self.maintain_state = False
  
          with warnings.catch_warnings(record=False):
              warnings.filterwarnings("ignore")
              # the following line issues a warning about the loss attribute being saved
              # but we do want to save it
              self.save_hyperparameters()  # Allows instantiation from a checkpoint from class
          self.random_seed = random_seed
          pl.seed_everything(self.random_seed, workers=True)
  
          # Loss
          self.loss = loss
          if valid_loss is None:
              self.valid_loss = loss
          else:
              self.valid_loss = valid_loss
          self.train_trajectories: List = []
          self.valid_trajectories: List = []
  
          # Optimization
          if optimizer is not None and not issubclass(optimizer, torch.optim.Optimizer):
              raise TypeError(
                  "optimizer is not a valid subclass of torch.optim.Optimizer"
              )
          self.optimizer = optimizer
          self.optimizer_kwargs = optimizer_kwargs if optimizer_kwargs is not None else {}
  
          # lr scheduler
          if lr_scheduler is not None and not issubclass(
              lr_scheduler, torch.optim.lr_scheduler.LRScheduler
          ):
              raise TypeError(
                  "lr_scheduler is not a valid subclass of torch.optim.lr_scheduler.LRScheduler"
              )
          self.lr_scheduler = lr_scheduler
          self.lr_scheduler_kwargs = (
              lr_scheduler_kwargs if lr_scheduler_kwargs is not None else {}
          )
  
          # Variables
          self.futr_exog_list = list(futr_exog_list) if futr_exog_list is not None else []
          self.hist_exog_list = list(hist_exog_list) if hist_exog_list is not None else []
          self.stat_exog_list = list(stat_exog_list) if stat_exog_list is not None else []
  
          # Set data sizes
          self.futr_exog_size = len(self.futr_exog_list)
          self.hist_exog_size = len(self.hist_exog_list)
          self.stat_exog_size = len(self.stat_exog_list)
  
          # Check if model supports exogenous, otherwise raise Exception
          if not self.EXOGENOUS_FUTR and self.futr_exog_size > 0:
              raise Exception(
                  f"{type(self).__name__} does not support future exogenous variables."
              )
          if not self.EXOGENOUS_HIST and self.hist_exog_size > 0:
              raise Exception(
                  f"{type(self).__name__} does not support historical exogenous variables."
              )
          if not self.EXOGENOUS_STAT and self.stat_exog_size > 0:
              raise Exception(
                  f"{type(self).__name__} does not support static exogenous variables."
              )
  
          # Protections for loss functions
          if isinstance(self.loss, (losses.IQLoss, losses.HuberIQLoss)):
              loss_type = type(self.loss)
              if not isinstance(self.valid_loss, loss_type):
                  raise Exception(
                      f"Please set valid_loss={type(self.loss).__name__}() when training with {type(self.loss).__name__}"
                  )
          if isinstance(self.loss, (losses.MQLoss, losses.HuberMQLoss)):
              if not isinstance(self.valid_loss, (losses.MQLoss, losses.HuberMQLoss)):
                  raise Exception(
                      f"Please set valid_loss to MQLoss() or HuberMQLoss() when training with {type(self.loss).__name__}"
                  )
          if isinstance(self.valid_loss, (losses.IQLoss, losses.HuberIQLoss)):
              valid_loss_type = type(self.valid_loss)
              if not isinstance(self.loss, valid_loss_type):
                  raise Exception(
                      f"Please set loss={type(self.valid_loss).__name__}() when validating with {type(self.valid_loss).__name__}"
                  )
  
          # Deny impossible loss / valid_loss combinations
          if (
              isinstance(self.loss, losses.BasePointLoss)
              and self.valid_loss.is_distribution_output
          ):
              raise Exception(
                  f"Validation with distribution loss {type(self.valid_loss).__name__} is not possible when using loss={type(self.loss).__name__}. Please use a point valid_loss (MAE, MSE, ...)"
              )
          elif self.valid_loss.is_distribution_output and self.valid_loss is not loss:
              # Maybe we should raise a Warning or an Exception here, but meh for now.
              self.valid_loss = loss
  
          if isinstance(self.loss, (losses.relMSE, losses.Accuracy, losses.sCRPS)):
              raise Exception(
                  f"{type(self.loss).__name__} cannot be used for training. Please use another loss function (MAE, MSE, ...)"
              )
  
          if isinstance(self.valid_loss, (losses.relMSE)):
              raise Exception(
                  f"{type(self.valid_loss).__name__} cannot be used for validation. Please use another valid_loss (MAE, MSE, ...)"
              )
  
          ## Trainer arguments ##
          # Max steps, validation steps and check_val_every_n_epoch
          trainer_kwargs = {**trainer_kwargs, "max_steps": max_steps}
  
          if "max_epochs" in trainer_kwargs.keys():
              raise Exception("max_epochs is deprecated, use max_steps instead.")
  
          # Callbacks
          if early_stop_patience_steps > 0:
              if "callbacks" not in trainer_kwargs:
                  trainer_kwargs["callbacks"] = []
              trainer_kwargs["callbacks"].append(
                  EarlyStopping(
                      monitor="ptl/val_loss", patience=early_stop_patience_steps
                  )
              )
  
          # Add GPU accelerator if available
          if trainer_kwargs.get("accelerator", None) is None:
              if torch.cuda.is_available():
                  trainer_kwargs["accelerator"] = "gpu"
          if trainer_kwargs.get("devices", None) is None:
              if torch.cuda.is_available():
                  trainer_kwargs["devices"] = -1
  
          # Avoid saturating local memory, disabled fit model checkpoints
          if trainer_kwargs.get("enable_checkpointing", None) is None:
              trainer_kwargs["enable_checkpointing"] = False
  
          # Set other attributes
          self.trainer_kwargs = trainer_kwargs
          self.h = h
          self.input_size = input_size
          self.windows_batch_size = windows_batch_size
          self.start_padding_enabled = start_padding_enabled
          self.predict_horizon = self.horizon_backup  # Used in recurrent prediction whereby predict h > h_train
  
          # Padder to complete train windows,
          # example y=[1,2,3,4,5] h=3 -> last y_output = [5,0,0]
          if start_padding_enabled:
              self.padder_train = nn.ConstantPad1d(
                  padding=(self.input_size - 1, self.h), value=0.0
              )
          else:
              self.padder_train = nn.ConstantPad1d(padding=(0, self.h), value=0.0)
  
          # Batch sizes
          if self.MULTIVARIATE and n_series is not None:
              self.batch_size = max(batch_size, n_series)
              if valid_batch_size is not None:
                  valid_batch_size = max(valid_batch_size, n_series)
          else:
              self.batch_size = batch_size
  
          if valid_batch_size is None:
              self.valid_batch_size = self.batch_size
          else:
              self.valid_batch_size = valid_batch_size
  
          if inference_windows_batch_size is None:
              self.inference_windows_batch_size = windows_batch_size
          else:
              self.inference_windows_batch_size = inference_windows_batch_size
  
          # Filtering training windows by available sample fractions
          if isinstance(training_data_availability_threshold, int):
              raise ValueError(
                  "training_data_availability_threshold cannot be an integer - must be a float"
              )
          elif isinstance(training_data_availability_threshold, float):
              if (
                  training_data_availability_threshold < 0.0
                  or training_data_availability_threshold > 1.0
              ):
                  raise ValueError(
                      f"training_data_availability_threshold must be between 0.0 and 1.0, got {training_data_availability_threshold}"
                  )
              self.min_insample_fraction = training_data_availability_threshold
              self.min_outsample_fraction = training_data_availability_threshold
          elif (
              isinstance(training_data_availability_threshold, (list, tuple))
              and len(training_data_availability_threshold) == 2
          ):
              for i, value in enumerate(training_data_availability_threshold):
                  if isinstance(value, int):
                      raise ValueError(
                          f"training_data_availability_threshold[{i}] cannot be an integer - must be a float"
                      )
                  if not isinstance(value, float):
                      raise ValueError(
                          f"training_data_availability_threshold[{i}] must be a float"
                      )
                  if value < 0.0 or value > 1.0:
                      raise ValueError(
                          f"training_data_availability_threshold[{i}] must be between 0.0 and 1.0, got {value}"
                      )
  
              self.min_insample_fraction = training_data_availability_threshold[0]
              self.min_outsample_fraction = training_data_availability_threshold[1]
          else:
              raise ValueError(
                  "training_data_availability_threshold must be a float or a list/tuple of two floats"
              )
  
          # Optimization
          self.learning_rate = learning_rate
          self.max_steps = max_steps
          self.num_lr_decays = num_lr_decays
          self.lr_decay_steps = (
              max(max_steps // self.num_lr_decays, 1) if self.num_lr_decays > 0 else 10e7
          )
          self.early_stop_patience_steps = early_stop_patience_steps
          self.val_check_steps = val_check_steps
          self.windows_batch_size = windows_batch_size
          self.step_size = step_size
  
          # If the model does not support exogenous, it can't support exclude_insample_y
          if exclude_insample_y and not (
              self.EXOGENOUS_FUTR or self.EXOGENOUS_HIST or self.EXOGENOUS_STAT
          ):
              raise Exception(
                  f"{type(self).__name__} does not support `exclude_insample_y=True`. Please set `exclude_insample_y=False`"
              )
  
          self.exclude_insample_y = exclude_insample_y
  
          # Scaler
          self.scaler = TemporalNorm(
              scaler_type=scaler_type,
              dim=1,  # Time dimension is 1.
              num_features=1 + len(self.hist_exog_list) + len(self.futr_exog_list),
          )
  
          # Fit arguments
          self.val_size = 0
          self.test_size = 0
  
          # Model state
          self.decompose_forecast = False
  
          # DataModule arguments
          self.dataloader_kwargs = dataloader_kwargs
          self.drop_last_loader = drop_last_loader
          # used by on_validation_epoch_end hook
          self.validation_step_outputs: List = []
          self.alias = alias
  
      def __repr__(self):
          return type(self).__name__ if self.alias is None else self.alias
  
      def _check_exog(self, dataset):
          temporal_cols = set(dataset.temporal_cols.tolist())
          static_cols = set(
              dataset.static_cols.tolist() if dataset.static_cols is not None else []
          )
  
          missing_hist = set(self.hist_exog_list) - temporal_cols
          missing_futr = set(self.futr_exog_list) - temporal_cols
          missing_stat = set(self.stat_exog_list) - static_cols
          if missing_hist:
              raise Exception(
                  f"{missing_hist} historical exogenous variables not found in input dataset"
              )
          if missing_futr:
              raise Exception(
                  f"{missing_futr} future exogenous variables not found in input dataset"
              )
          if missing_stat:
              raise Exception(
                  f"{missing_stat} static exogenous variables not found in input dataset"
              )
  
      def _restart_seed(self, random_seed):
          if random_seed is None:
              random_seed = self.random_seed
          torch.manual_seed(random_seed)
  
      def _get_temporal_exogenous_cols(self, temporal_cols):
          return list(
              set(temporal_cols.tolist()) & set(self.hist_exog_list + self.futr_exog_list)
          )
  
      def _set_quantiles(self, quantiles=None):
          if quantiles is None and isinstance(
              self.loss, (losses.IQLoss, losses.HuberIQLoss)
          ):
              self.loss.update_quantile(q=[0.5])
          elif hasattr(self.loss, "update_quantile") and callable(
              self.loss.update_quantile
          ):
              self.loss.update_quantile(q=quantiles)
  
      def _fit_distributed(
          self,
          distributed_config,
          datamodule,
          val_size,
          test_size,
      ):
          assert distributed_config is not None
          from pyspark.ml.torch.distributor import TorchDistributor
  
          def train_fn(
              model_cls,
              model_params,
              datamodule,
              trainer_kwargs,
              num_tasks,
              num_proc_per_task,
              val_size,
              test_size,
          ):
              import pytorch_lightning as pl
  
              # we instantiate here to avoid pickling large tensors (weights)
              model = model_cls(**model_params)
              model.val_size = val_size
              model.test_size = test_size
              for arg in ("devices", "num_nodes"):
                  trainer_kwargs.pop(arg, None)
              trainer = pl.Trainer(
                  strategy="ddp",
                  use_distributed_sampler=False,  # to ensure our dataloaders are used as-is
                  num_nodes=num_tasks,
                  devices=num_proc_per_task,
                  **trainer_kwargs,
              )
              trainer.fit(model=model, datamodule=datamodule)
              model.metrics = trainer.callback_metrics
              model.__dict__.pop("_trainer", None)
              return model
  
          def is_gpu_accelerator(accelerator):
              from pytorch_lightning.accelerators.cuda import CUDAAccelerator
  
              return (
                  accelerator == "gpu"
                  or isinstance(accelerator, CUDAAccelerator)
                  or (accelerator == "auto" and CUDAAccelerator.is_available())
              )
  
          local_mode = distributed_config.num_nodes == 1
          if local_mode:
              num_tasks = 1
              num_proc_per_task = distributed_config.devices
          else:
              num_tasks = distributed_config.num_nodes * distributed_config.devices
              num_proc_per_task = 1  # number of GPUs per task
          num_proc = num_tasks * num_proc_per_task
          use_gpu = is_gpu_accelerator(self.trainer_kwargs["accelerator"])
          model = TorchDistributor(
              num_processes=num_proc,
              local_mode=local_mode,
              use_gpu=use_gpu,
          ).run(
              train_fn,
              model_cls=type(self),
              model_params=self.hparams,
              datamodule=datamodule,
              trainer_kwargs=self.trainer_kwargs,
              num_tasks=num_tasks,
              num_proc_per_task=num_proc_per_task,
              val_size=val_size,
              test_size=test_size,
          )
          return model
  
      def _fit(
          self,
          dataset,
          batch_size,
          valid_batch_size=1024,
          val_size=0,
          test_size=0,
          random_seed=None,
          shuffle_train=True,
          distributed_config=None,
      ):
          self._check_exog(dataset)
          self._restart_seed(random_seed)
  
          self.val_size = val_size
          self.test_size = test_size
          is_local = isinstance(dataset, BaseTimeSeriesDataset)
          if is_local:
              datamodule_constructor = TimeSeriesDataModule
          else:
              datamodule_constructor = _DistributedTimeSeriesDataModule
  
          dataloader_kwargs = (
              self.dataloader_kwargs if self.dataloader_kwargs is not None else {}
          )
          datamodule = datamodule_constructor(
              dataset=dataset,
              batch_size=batch_size,
              valid_batch_size=valid_batch_size,
              drop_last=self.drop_last_loader,
              shuffle_train=shuffle_train,
              **dataloader_kwargs,
          )
  
          if self.val_check_steps > self.max_steps:
              warnings.warn(
                  "val_check_steps is greater than max_steps, "
                  "setting val_check_steps to max_steps."
              )
          val_check_interval = min(self.val_check_steps, self.max_steps)
          self.trainer_kwargs["val_check_interval"] = int(val_check_interval)
          self.trainer_kwargs["check_val_every_n_epoch"] = None
  
          if is_local:
              model = self
              trainer = pl.Trainer(**model.trainer_kwargs)
              trainer.fit(model, datamodule=datamodule)
              model.metrics = trainer.callback_metrics
              model.__dict__.pop("_trainer", None)
          else:
              model = self._fit_distributed(
                  distributed_config,
                  datamodule,
                  val_size,
                  test_size,
              )
          return model
  
      def on_fit_start(self):
          torch.manual_seed(self.random_seed)
          np.random.seed(self.random_seed)
          random.seed(self.random_seed)
  
      def configure_optimizers(self):
          if self.optimizer:
              optimizer_signature = inspect.signature(self.optimizer)
              optimizer_kwargs = deepcopy(self.optimizer_kwargs)
              if "lr" in optimizer_signature.parameters:
                  if "lr" in optimizer_kwargs:
                      warnings.warn(
                          "ignoring learning rate passed in optimizer_kwargs, using the model's learning rate"
                      )
                  optimizer_kwargs["lr"] = self.learning_rate
              optimizer = self.optimizer(params=self.parameters(), **optimizer_kwargs)
          else:
              if self.optimizer_kwargs:
                  warnings.warn(
                      "ignoring optimizer_kwargs as the optimizer is not specified"
                  )
              optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
  
          lr_scheduler = {"frequency": 1, "interval": "step"}
          if self.lr_scheduler:
              lr_scheduler_signature = inspect.signature(self.lr_scheduler)
              lr_scheduler_kwargs = deepcopy(self.lr_scheduler_kwargs)
              if "optimizer" in lr_scheduler_signature.parameters:
                  if "optimizer" in lr_scheduler_kwargs:
                      warnings.warn(
                          "ignoring optimizer passed in lr_scheduler_kwargs, using the model's optimizer"
                      )
                      del lr_scheduler_kwargs["optimizer"]
              lr_scheduler["scheduler"] = self.lr_scheduler(
                  optimizer=optimizer, **lr_scheduler_kwargs
              )
          else:
              if self.lr_scheduler_kwargs:
                  warnings.warn(
                      "ignoring lr_scheduler_kwargs as the lr_scheduler is not specified"
                  )
              lr_scheduler["scheduler"] = torch.optim.lr_scheduler.StepLR(
                  optimizer=optimizer, step_size=self.lr_decay_steps, gamma=0.5
              )
          return {"optimizer": optimizer, "lr_scheduler": lr_scheduler}
  
      def get_test_size(self):
          return self.test_size
  
      def set_test_size(self, test_size):
          self.test_size = test_size
  
      def on_validation_epoch_end(self):
          if self.val_size == 0:
              return
          losses = torch.stack(self.validation_step_outputs)
          avg_loss = losses.mean().detach().item()
          self.log(
              "ptl/val_loss",
              avg_loss,
              batch_size=losses.size(0),
              sync_dist=True,
          )
          self.valid_trajectories.append((self.global_step, avg_loss))
          self.validation_step_outputs.clear()  # free memory (compute `avg_loss` per epoch)
  
      def save(self, path):
          with fsspec.open(path, "wb") as f:
              torch.save(
                  {"hyper_parameters": self.hparams, "state_dict": self.state_dict()},
                  f,
              )
  
      @classmethod
      def load(cls, path, **kwargs):
          if "weights_only" in inspect.signature(torch.load).parameters:
              kwargs["weights_only"] = False
          with fsspec.open(path, "rb") as f, warnings.catch_warnings():
              # ignore possible warnings about weights_only=False
              warnings.filterwarnings("ignore", category=FutureWarning)
              content = torch.load(f, **kwargs)
          with _disable_torch_init():
              model = cls(**content["hyper_parameters"])
          if "assign" in inspect.signature(model.load_state_dict).parameters:
              model.load_state_dict(content["state_dict"], strict=True, assign=True)
          else:  # pytorch<2.1
              model.load_state_dict(content["state_dict"], strict=True)
          return model
  
      def _create_windows(self, batch, step):
          # Parse common data
          window_size = self.input_size + self.h
          temporal_cols = batch["temporal_cols"]
          temporal = batch["temporal"]
  
          if step == "train":
              if self.val_size + self.test_size > 0:
                  cutoff = -self.val_size - self.test_size
                  temporal = temporal[:, :, :cutoff]
  
              temporal = self.padder_train(temporal)
  
              if temporal.shape[-1] < window_size:
                  raise Exception(
                      "Time series is too short for training, consider setting a smaller input size or set start_padding_enabled=True"
                  )
  
              windows = temporal.unfold(
                  dimension=-1, size=window_size, step=self.step_size
              )
  
              if self.MULTIVARIATE:
                  # [n_series, C, Ws, L + h] -> [Ws, L + h, C, n_series]
                  windows = windows.permute(2, 3, 1, 0)
              else:
                  # [n_series, C, Ws, L + h] -> [Ws * n_series, L + h, C, 1]
                  windows_per_serie = windows.shape[2]
                  windows = windows.permute(0, 2, 3, 1)
                  windows = windows.flatten(0, 1)
                  windows = windows.unsqueeze(-1)
  
              # Calculate minimum required available points based on fractions
              min_insample_points = max(
                  1, int(self.input_size * self.min_insample_fraction * self.n_series)
              )
              min_outsample_points = max(
                  1, int(self.h * self.min_outsample_fraction * self.n_series)
              )
  
              # Sample based on available conditions
              available_idx = temporal_cols.get_loc("available_mask")
              insample_condition = windows[:, : self.input_size, available_idx]
              insample_condition = torch.sum(
                  insample_condition, axis=(1, -1)
              )  # Sum over time & series dimension
              final_condition = insample_condition >= min_insample_points
  
              if self.h > 0:
                  outsample_condition = windows[:, self.input_size :, available_idx]
                  outsample_condition = torch.sum(
                      outsample_condition, axis=(1, -1)
                  )  # Sum over time & series dimension
                  final_condition = (outsample_condition >= min_outsample_points) & (
                      insample_condition >= min_insample_points
                  )
  
              windows = windows[final_condition]
  
              # Parse Static data to match windows
              static = batch.get("static", None)
              static_cols = batch.get("static_cols", None)
  
              # Repeat static if univariate: [n_series, S] -> [Ws * n_series, S]
              if static is not None and not self.MULTIVARIATE:
                  static = torch.repeat_interleave(
                      static, repeats=windows_per_serie, dim=0
                  )
                  static = static[final_condition]
  
              # Protection of empty windows
              if final_condition.sum() == 0:
                  raise Exception("No windows available for training")
  
              return windows, static, static_cols
  
          elif step in ["predict", "val"]:
  
              if step == "predict":
                  initial_input = temporal.shape[-1] - self.test_size
                  if (
                      initial_input <= self.input_size
                  ):  # There is not enough data to predict first timestamp
                      temporal = F.pad(
                          temporal,
                          pad=(self.input_size - initial_input, 0),
                          mode="constant",
                          value=0.0,
                      )
                  predict_step_size = self.predict_step_size
                  cutoff = -self.input_size - self.test_size
                  temporal = temporal[:, :, cutoff:]
  
              elif step == "val":
                  predict_step_size = self.step_size
                  cutoff = -self.input_size - self.val_size - self.test_size
                  if self.test_size > 0:
                      temporal = batch["temporal"][:, :, cutoff : -self.test_size]
                  else:
                      temporal = batch["temporal"][:, :, cutoff:]
                  if temporal.shape[-1] < window_size:
                      initial_input = temporal.shape[-1] - self.val_size
                      temporal = F.pad(
                          temporal,
                          pad=(self.input_size - initial_input, 0),
                          mode="constant",
                          value=0.0,
                      )
  
              if (
                  (step == "predict")
                  and (self.test_size == 0)
                  and (len(self.futr_exog_list) == 0)
              ):
                  temporal = F.pad(temporal, pad=(0, self.h), mode="constant", value=0.0)
  
              windows = temporal.unfold(
                  dimension=-1, size=window_size, step=predict_step_size
              )
  
              static = batch.get("static", None)
              static_cols = batch.get("static_cols", None)
  
              if self.MULTIVARIATE:
                  # [n_series, C, Ws, L + h] -> [Ws, L + h, C, n_series]
                  windows = windows.permute(2, 3, 1, 0)
              else:
                  # [n_series, C, Ws, L + h] -> [Ws * n_series, L + h, C, 1]
                  windows_per_serie = windows.shape[2]
                  windows = windows.permute(0, 2, 3, 1)
                  windows = windows.flatten(0, 1)
                  windows = windows.unsqueeze(-1)
                  if static is not None:
                      static = torch.repeat_interleave(
  