docstr: PyTorch module implementing the N-BEATS architecture.
  
  Parameters
  ----------
  output_dim
      Number of output components in the target
  nr_params
      The number of parameters of the likelihood (or 1 if no likelihood is used).
  generic_architecture
      Boolean value indicating whether the generic architecture of N-BEATS is used.
      If not, the interpretable architecture outlined in the paper (consisting of one trend
      and one seasonality stack with appropriate waveform generator functions).
  num_stacks
      The number of stacks that make up the whole model. Only used if `generic_architecture` is set to `True`.
  num_blocks
      The number of blocks making up every stack.
  num_layers
      The number of fully connected layers preceding the final forking layers in each block of every stack.
  layer_widths
      Determines the number of neurons that make up each fully connected layer in each block of every stack.
      If a list is passed, it must have a length equal to `num_stacks` and every entry in that list corresponds
      to the layer width of the corresponding stack. If an integer is passed, every stack will have blocks
      with FC layers of the same width.
  expansion_coefficient_dim
      The dimensionality of the waveform generator parameters, also known as expansion coefficients.
      Only used if `generic_architecture` is set to `True`.
  trend_polynomial_degree
      The degree of the polynomial used as waveform generator in trend stacks. Only used if
      `generic_architecture` is set to `False`.
  batch_norm
      Whether to apply batch norm on first block of the first stack
  dropout
      Dropout probability
  activation
      The activation function of encoder/decoder intermediate layer.
  **kwargs
      all parameters required for :class:`darts.model.forecasting_models.PLForecastingModule` base class.
  
  Inputs
  ------
  x of shape `(batch_size, input_chunk_length)`
      Tensor containing the input sequence.
  
  Outputs
  -------
  y of shape `(batch_size, output_chunk_length, target_size/output_dim, nr_params)`
      Tensor containing the output of the NBEATS module.
docstr_tok: ["PyTorch","module","implementing","the","N-BEATS","architecture",".","Parameters","--","--","--","--","--","output_dim","Number","of","output","components","in","the","target","nr_params","The","number","of","parameters","of","the","likelihood","(","or","1","if","no","likelihood","is","used",")",".","generic_architecture","Boolean","value","indicating","whether","the","generic","architecture","of","N-BEATS","is","used",".","If","not",",","the","interpretable","architecture","outlined","in","the","paper","(","consisting","of","one","trend","and","one","seasonality","stack","with","appropriate","waveform","generator","functions",")",".","num_stacks","The","number","of","stacks","that","make","up","the","whole","model",".","Only","used","if","`","generic_architecture","`","is","set","to","`","True","`",".","num_blocks","The","number","of","blocks","making","up","every","stack",".","num_layers","The","number","of","fully","connected","layers","preceding","the","final","forking","layers","in","each","block","of","every","stack",".","layer_widths","Determines","the","number","of","neurons","that","make","up","each","fully","connected","layer","in","each","block","of","every","stack",".","If","a","list","is","passed",",","it","must","have","a","length","equal","to","`","num_stacks","`","and","every","entry","in","that","list","corresponds","to","the","layer","width","of","the","corresponding","stack",".","If","an","integer","is","passed",",","every","stack","will","have","blocks","with","FC","layers","of","the","same","width",".","expansion_coefficient_dim","The","dimensionality","of","the","waveform","generator","parameters",",","also","known","as","expansion","coefficients",".","Only","used","if","`","generic_architecture","`","is","set","to","`","True","`",".","trend_polynomial_degree","The","degree","of","the","polynomial","used","as","waveform","generator","in","trend","stacks",".","Only","used","if","`","generic_architecture","`","is","set","to","`","False","`",".","batch_norm","Whether","to","apply","batch","norm","on","first","block","of","the","first","stack","dropout","Dropout","probability","activation","The","activation","function","of","encoder\/decoder","intermediate","layer",".","*","*","kwargs","all","parameters","required","for",":","class",":","`","darts.model.forecasting_models.PLForecastingModule","`","base","class",".","Inputs","--","--","--","x","of","shape","`","(","batch_size",",","input_chunk_length",")","`","Tensor","containing","the","input","sequence",".","Outputs","--","--","--","-","y","of","shape","`","(","batch_size",",","output_chunk_length",",","target_size\/output_dim",",","nr_params",")","`","Tensor","containing","the","output","of","the","NBEATS","module","."]
code: def __init__(
      self,
      input_dim: int,
      output_dim: int,
      nr_params: int,
      generic_architecture: bool,
      num_stacks: int,
      num_blocks: int,
      num_layers: int,
      layer_widths: List[int],
      expansion_coefficient_dim: int,
      trend_polynomial_degree: int,
      batch_norm: bool,
      dropout: float,
      activation: str,
      **kwargs,
  ):
      """PyTorch module implementing the N-BEATS architecture.
  
      Parameters
      ----------
      output_dim
          Number of output components in the target
      nr_params
          The number of parameters of the likelihood (or 1 if no likelihood is used).
      generic_architecture
          Boolean value indicating whether the generic architecture of N-BEATS is used.
          If not, the interpretable architecture outlined in the paper (consisting of one trend
          and one seasonality stack with appropriate waveform generator functions).
      num_stacks
          The number of stacks that make up the whole model. Only used if `generic_architecture` is set to `True`.
      num_blocks
          The number of blocks making up every stack.
      num_layers
          The number of fully connected layers preceding the final forking layers in each block of every stack.
      layer_widths
          Determines the number of neurons that make up each fully connected layer in each block of every stack.
          If a list is passed, it must have a length equal to `num_stacks` and every entry in that list corresponds
          to the layer width of the corresponding stack. If an integer is passed, every stack will have blocks
          with FC layers of the same width.
      expansion_coefficient_dim
          The dimensionality of the waveform generator parameters, also known as expansion coefficients.
          Only used if `generic_architecture` is set to `True`.
      trend_polynomial_degree
          The degree of the polynomial used as waveform generator in trend stacks. Only used if
          `generic_architecture` is set to `False`.
      batch_norm
          Whether to apply batch norm on first block of the first stack
      dropout
          Dropout probability
      activation
          The activation function of encoder/decoder intermediate layer.
      **kwargs
          all parameters required for :class:`darts.model.forecasting_models.PLForecastingModule` base class.
  
      Inputs
      ------
      x of shape `(batch_size, input_chunk_length)`
          Tensor containing the input sequence.
  
      Outputs
      -------
      y of shape `(batch_size, output_chunk_length, target_size/output_dim, nr_params)`
          Tensor containing the output of the NBEATS module.
  
      """
      super().__init__(**kwargs)
  
      self.input_dim = input_dim
      self.output_dim = output_dim
      self.nr_params = nr_params
      self.input_chunk_length_multi = (
          self.input_chunk_length * input_dim
      )
      self.target_length = self.output_chunk_length * input_dim
      self.dropout = dropout
      self.batch_norm = batch_norm
      self.activation = activation
  
      if generic_architecture:
          self.stacks_list = [
              _Stack(
                  num_blocks,
                  num_layers,
                  layer_widths[i],
                  nr_params,
                  expansion_coefficient_dim,
                  self.input_chunk_length_multi,
                  self.target_length,
                  _GType.GENERIC,
                  batch_norm=(
                      self.batch_norm and i == 0
                  ),  # batch norm only on first block of first stack
                  dropout=self.dropout,
                  activation=self.activation,
              )
              for i in range(num_stacks)
          ]
      else:
          num_stacks = 2
          trend_stack = _Stack(
              num_blocks,
              num_layers,
              layer_widths[0],
              nr_params,
              trend_polynomial_degree + 1,
              self.input_chunk_length_multi,
              self.target_length,
              _GType.TREND,
              batch_norm=self.batch_norm,
              dropout=self.dropout,
              activation=self.activation,
          )
          seasonality_stack = _Stack(
              num_blocks,
              num_layers,
              layer_widths[1],
              nr_params,
              -1,
              self.input_chunk_length_multi,
              self.target_length,
              _GType.SEASONALITY,
              batch_norm=self.batch_norm,
              dropout=self.dropout,
              activation=self.activation,
          )
          self.stacks_list = [trend_stack, seasonality_stack]
  
      self.stacks = nn.ModuleList(self.stacks_list)
  
      # setting the last backcast "branch" to be not trainable (without next block/stack, it doesn't need to be
      # backpropagated). Removing this lines would cause logtensorboard to crash, since no gradient is stored
      # on this params (the last block backcast is not part of the final output of the net).
      self.stacks_list[-1].blocks[
          -1
      ].backcast_linear_layer.requires_grad_(False)
      self.stacks_list[-1].blocks[-1].backcast_g.requires_grad_(False)
  
imports: from enum import Enum
  from typing import List, NewType, Tuple, Union
  import numpy
  import torch
  import torch.nn
  from darts.logging import get_logger, raise_if_not, raise_log
  from darts.models.forecasting.pl_forecasting_module import PLPastCovariatesModule, io_processor
  from darts.models.forecasting.torch_forecasting_model import PastCovariatesTorchModel
  from darts.utils.torch import MonteCarloDropout
label_desc: Use image node_id 1 to create a new _NBEATSModule object from inherited base classes: PLPastCovariatesModule with example: obj = _NBEATSModule(input_dim, output_dim, nr_params, generic_architecture, num_stacks, num_blocks, num_layers, layer_widths, expansion_coefficient_dim, trend_polynomial_degree, batch_norm, dropout, activation)