class_docstr: NHITS block which takes a basis function as an argument.
class_docstr_tok: ["NHITS","block","which","takes","a","basis","function","as","an","argument","."]
code: def __init__(
      self,
      input_size: int,
      h: int,
      n_theta: int,
      mlp_units: list,
      basis: nn.Module,
      futr_input_size: int,
      hist_input_size: int,
      stat_input_size: int,
      n_pool_kernel_size: int,
      pooling_mode: str,
      dropout_prob: float,
      activation: str,
  ):
      super().__init__()
  
      pooled_hist_size = int(np.ceil(input_size / n_pool_kernel_size))
      pooled_futr_size = int(
          np.ceil((input_size + h) / n_pool_kernel_size)
      )
  
      input_size = (
          pooled_hist_size
          + hist_input_size * pooled_hist_size
          + futr_input_size * pooled_futr_size
          + stat_input_size
      )
  
      self.dropout_prob = dropout_prob
      self.futr_input_size = futr_input_size
      self.hist_input_size = hist_input_size
      self.stat_input_size = stat_input_size
  
      assert (
          activation in ACTIVATIONS
      ), f"{activation} is not in {ACTIVATIONS}"
      assert (
          pooling_mode in POOLING
      ), f"{pooling_mode} is not in {POOLING}"
  
      activ = getattr(nn, activation)()
  
      self.pooling_layer = getattr(nn, pooling_mode)(
          kernel_size=n_pool_kernel_size,
          stride=n_pool_kernel_size,
          ceil_mode=True,
      )
  
      # Block MLPs
      hidden_layers = [
          nn.Linear(
              in_features=input_size, out_features=mlp_units[0][0]
          )
      ]
      for layer in mlp_units:
          hidden_layers.append(
              nn.Linear(in_features=layer[0], out_features=layer[1])
          )
          hidden_layers.append(activ)
  
          if self.dropout_prob > 0:
              # raise NotImplementedError('dropout')
              hidden_layers.append(nn.Dropout(p=self.dropout_prob))
  
      output_layer = [
          nn.Linear(in_features=mlp_units[-1][1], out_features=n_theta)
      ]
      layers = hidden_layers + output_layer
      self.layers = nn.Sequential(*layers)
      self.basis = basis
  
imports: from typing import Tuple, Optional
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from ..losses.pytorch import MAE
  from ..common._base_windows import BaseWindows
label_desc: Use image node_id 1 to create a new NHITSBlock object from inherited base classes: nn with example: obj = NHITSBlock(input_size, h, n_theta, mlp_units, basis, futr_input_size, hist_input_size, stat_input_size, n_pool_kernel_size, pooling_mode, dropout_prob, activation)