class_docstr: N-HiTS block which takes a basis function as an argument.
class_docstr_tok: ["N-HiTS","block","which","takes","a","basis","function","as","an","argument","."]
code: def __init__(
      self,
      context_length: int,
      prediction_length: int,
      output_size: int,
      encoder_covariate_size: int,
      decoder_covariate_size: int,
      static_size: int,
      static_hidden_size: int,
      n_theta: int,
      hidden_size: List[int],
      pooling_sizes: int,
      pooling_mode: str,
      basis: nn.Module,
      n_layers: int,
      batch_normalization: bool,
      dropout: float,
      activation: str,
  ):
      super().__init__()
  
      assert pooling_mode in ["max", "average"]
  
      self.context_length_pooled = int(
          np.ceil(context_length / pooling_sizes)
      )
  
      if static_size == 0:
          static_hidden_size = 0
  
      self.context_length = context_length
      self.output_size = output_size
      self.n_theta = n_theta
      self.prediction_length = prediction_length
      self.static_size = static_size
      self.static_hidden_size = static_hidden_size
      self.encoder_covariate_size = encoder_covariate_size
      self.decoder_covariate_size = decoder_covariate_size
      self.pooling_sizes = pooling_sizes
      self.batch_normalization = batch_normalization
      self.dropout = dropout
  
      self.hidden_size = [
          self.context_length_pooled * len(self.output_size)
          + self.context_length * self.encoder_covariate_size
          + self.prediction_length * self.decoder_covariate_size
          + self.static_hidden_size
      ] + hidden_size
  
      assert (
          activation in ACTIVATIONS
      ), f"{activation} is not in {ACTIVATIONS}"
      activ = getattr(nn, activation)()
  
      if pooling_mode == "max":
          self.pooling_layer = nn.MaxPool1d(
              kernel_size=self.pooling_sizes,
              stride=self.pooling_sizes,
              ceil_mode=True,
          )
      elif pooling_mode == "average":
          self.pooling_layer = nn.AvgPool1d(
              kernel_size=self.pooling_sizes,
              stride=self.pooling_sizes,
              ceil_mode=True,
          )
  
      hidden_layers = []
      for i in range(n_layers):
          hidden_layers.append(
              nn.Linear(
                  in_features=self.hidden_size[i],
                  out_features=self.hidden_size[i + 1],
              )
          )
          hidden_layers.append(activ)
  
          if self.batch_normalization:
              hidden_layers.append(
                  nn.BatchNorm1d(num_features=self.hidden_size[i + 1])
              )
  
          if self.dropout > 0:
              hidden_layers.append(nn.Dropout(p=self.dropout))
  
      output_layer = [
          nn.Linear(
              in_features=self.hidden_size[-1],
              out_features=context_length * len(output_size)
              + n_theta * sum(output_size),
          )
      ]
      layers = hidden_layers + output_layer
  
      # static_size is computed with data, static_hidden_size is provided by user, if 0 no statics are used
      if (self.static_size > 0) and (self.static_hidden_size > 0):
          self.static_encoder = StaticFeaturesEncoder(
              in_features=static_size, out_features=static_hidden_size
          )
      self.layers = nn.Sequential(*layers)
      self.basis = basis
  
imports: from functools import partial
  from typing import List, Tuple
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new NHiTSBlock object from inherited base classes: nn with example: obj = NHiTSBlock(context_length, prediction_length, output_size, encoder_covariate_size, decoder_covariate_size, static_size, static_hidden_size, n_theta, hidden_size, pooling_sizes, pooling_mode, basis, n_layers, batch_normalization, dropout, activation)