text: # AUTOGENERATED! DO NOT EDIT! File to edit: blocks.ipynb (unless otherwise specified).
  
  __all__ = ['DefaultModel', 'DefaultQModel']
  
  # Cell
  import numpy as np
  import matplotlib.pyplot as plt
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from ..data import convert_date
  from .fourier import Seasonal
  from .trend import Trend
  from .squasher import Squasher
  from .linearX import LinearX
  
  torch.Tensor.ndim = property(lambda x: x.dim())
  
  # Cell
  class DefaultModel(nn.Module):
      """
      Sum of Linear Trend, Seasonality and squashed.
      """
      def __init__(self, moments, breakpoints=None, y_n=7, m_n=5, w_n=0, l=None, h=None):
          super().__init__()
          if 'x' in moments:
              dims = moments['x'][0].shape[1]
          else:
              dims = 0
  
          self.trend = Trend(breakpoints, moments)
          self.seasonal = Seasonal(y_n, m_n, w_n, scale=moments['t'][1])
          self.linear = LinearX(dims)
          self.squash = Squasher(l, h, *moments['y'])
  
      def forward(self, t, x=None):
          prediction = self.seasonal(t) + self.trend(t) + self.linear(x)
          prediction = self.squash(prediction)
          return prediction
  
  class DefaultQModel(nn.Module):
      """
      Same as DefaultModel but with multiple outputs corresponding to quantiles
      """
      def __init__(self, moments, breakpoints=None, y_n=7, m_n=5, w_n=0, l=None, h=None, quantiles=[0.05, 0.5, 0.95]):
          super().__init__()
          assert 0.5 in quantiles, f'0.5 needs to be in quantiles. Provided {quantiles} as quantiles.'
          self.idx = quantiles.index(0.5)
          signs = [q-0.5 for q in quantiles]
          self.signs = torch.Tensor([-1 if s<0 else 1 for i,s in enumerate(signs) if i != self.idx])[None,:]
          self.idxs = [i for i in range(len(quantiles)) if i != self.idx]
  
          median_args = {'y_n': y_n, 'm_n': m_n, 'w_n': w_n}
          other_args = {'y_n': 0, 'm_n': 0, 'w_n': 0}
          args = [median_args if q==0.5 else other_args for q in quantiles]
  
          self.models = nn.ModuleList([DefaultModel(moments, breakpoints, **arg) for arg in args])
          self.squash = Squasher(l, h, *moments['y'])
  
      def forward(self, t, x=None):
          prediction = torch.cat([m(t,x) for m in self.models], -1)
          median = prediction[:, [self.idx]]
          prediction[:, self.idxs] = median + F.softplus(prediction[:, self.idxs]) * self.signs
  
          return self.squash(prediction)