text: import torch
  import torch.nn.functional as F
  import numpy as np
  
  def mse(pred, true):
    return torch.mean((pred-true)**2)
  
  def rmse(pred, true):
    return torch.sqrt(mse(pred, true))
  
  def nig_nll(true, gamma, v, alpha, beta):
    
    twoBlambda = 2*beta*(1+v)
  
    nll = 0.5 * torch.log(torch.tensor(np.pi, dtype=torch.float32)/v) \
          - alpha*torch.log(twoBlambda) \
          + (alpha+0.5) * torch.log(v*(true-gamma)**2 + twoBlambda) \
          + torch.lgamma(alpha) \
          - torch.lgamma(alpha+0.5)
    return torch.mean(nll)
  
  def nig_reg(true, gamma, v, alpha, beta):
    
    error = torch.abs(true-gamma)
    evi = 2*v+alpha
    reg = error*evi
    return torch.mean(reg)
  
  def nig_loss(true, gamma, v, alpha, beta, coeff=1.0):
    loss_nll = nig_nll(true, gamma, v, alpha, beta)
    loss_reg = nig_reg(true, gamma, v, alpha, beta)
    loss =  loss_nll + coeff*loss_reg
    return loss, loss_nll, loss_reg
  
  
  def nll_loss(pred, true, sigma):
    return torch.mean((((pred-true)**2)/torch.exp(sigma))+sigma)
  
  class EvidentialNet(torch.nn.Module):
      epochs = 40
      batch_size = 8
      custom_name = 'Evidential'
      def __init__(self, **kwargs):
          super(EvidentialNet, self).__init__()
          n_feature = kwargs['n_features']
          n_hidden = kwargs['n_hidden']
          n_output = 1
          self.hidden_1 = torch.nn.Linear(n_feature, n_hidden)
          self.hidden_2 = torch.nn.Linear(n_hidden, n_hidden)
          self.gamma = torch.nn.Linear(n_hidden, n_output)   # mean prediction node
          self.softplus = torch.nn.Softplus()
          self.v = torch.nn.Linear(n_hidden, n_output)
          self.alpha = torch.nn.Linear(n_hidden, n_output)
          self.beta = torch.nn.Linear(n_hidden, n_output)
          self.coeff = 0.0
          self.epsilon = 1e-2
          self.maxi_rate = 1e-4
  
      def forward(self, x):
          x = F.relu(self.hidden_1(x))      # activation function for hidden layer
          x = self.hidden_2(x)
          gamma = self.gamma(x)             # linear output
          v = self.softplus(self.v(x))
          alpha = self.softplus(self.alpha(x))+1
          beta = self.softplus(self.beta(x))
  
          return gamma, v, alpha, beta
  
      def custom_compile(self):
          self.optimizer = torch.optim.Adam(self.parameters(), lr=5e-4)
          scheduler = None
          # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999)
          return self.optimizer, scheduler
  
      def train_step(self, **kwargs):
          self.train()
          gamma, upsilon, alpha, beta = self(kwargs['batch'][0])
          optim_m = kwargs['optim']
          # scheduler_m = kwargs['scheduler']      
          optim_m.zero_grad()
          loss, _, reg_l = nig_loss(kwargs['batch'][1], gamma, upsilon, alpha, beta, coeff=self.coeff)
          loss.backward()
          optim_m.step()
          self.coeff = self.coeff + (self.maxi_rate*(reg_l.item() - self.epsilon))
          # scheduler_m.step()
          # print(loss.item(), gamma, upsilon, alpha, beta)
          return loss.item()
  
      def evaluation_step(self,**kwargs):
          self.eval()
          target_scale = kwargs['target_scale']
          test_preds, _, _, _ = self(kwargs['test_tuple'][0])
          test_preds = test_preds
          test_score = kwargs['eval_func'](test_preds, kwargs['test_tuple'][1])
          scores = []
          for i in test_score:
            i = i*target_scale
            scores.append(i.detach().cpu().numpy())
  
          return scores
  
      def predict_step(self, **kwargs):
          self.eval()
          gamma, upsilon, alpha, beta = self(kwargs['batch'][0])
          gamma_vals = gamma.data.numpy()
          upsilon_vals = upsilon.data.numpy()
          alpha_vals = alpha.data.numpy()
          beta_vals = beta.data.numpy()
  
          mu = gamma_vals.squeeze()
          alea_unc = np.sqrt(beta_vals / (upsilon_vals*(alpha_vals -1)))
          alea_unc = np.minimum(alea_unc, 1e2)[:, 0]
  
          epi_unc = beta_vals / (alpha_vals -1)
          epi_unc = np.minimum(epi_unc, 1e2)[:, 0]
          return mu, alea_unc, epi_unc
  
  class MCDropout(torch.nn.Module):
      custom_name = "MCDropout"
      epochs = 40
      batch_size = 128
      def __init__(self, **kwargs):
          super(MCDropout, self).__init__()   
          n_features = kwargs['n_features']
          n_hidden = kwargs['n_hidden']
          tau = kwargs['tau']
          n_train = kwargs['n_train']   
          dropout_rate = 0.1
          lengthscale = 1e-2    
          lengthscale = 1e-2 
          reg = lengthscale**2 * (1 - dropout_rate) / (2. * n_train * tau * 1.4)
          self.reg = reg
          self.dropout1 = torch.nn.Dropout(dropout_rate)
          self.dropout2 = torch.nn.Dropout(dropout_rate)
          self.dropout3 = torch.nn.Dropout(dropout_rate)
          self.hidden = torch.nn.Linear(n_features, n_hidden)   # hidden layer
          self.hidden2 = torch.nn.Linear(n_hidden, n_hidden)
          self.hidden3 = torch.nn.Linear(n_hidden, n_hidden)
          self.predict = torch.nn.Linear(n_hidden, 1)   # output layer
          self.device = kwargs['device']
          self.relu = torch.nn.ReLU(inplace=True)
  
      def forward(self, x):
          x = self.relu(self.dropout1(self.hidden(x)))    # activation function for hidden layer
          x = self.relu(self.dropout2(self.hidden2(x)))
          # x = self.relu(self.dropout3(self.hidden3(x)))
          x = self.predict(x)             # linear output
          return x[:,0]
  
      def custom_compile(self):
        self.optimizer = torch.optim.Adam(self.parameters(), lr=0.01, weight_decay=self.reg)
        self.criterion = torch.nn.MSELoss(size_average=False)
        scheduler = None
        # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999)
        return self.optimizer, scheduler
  
      def train_step(self,**kwargs):
        kwargs['model'].train()
        preds = self(kwargs['batch'][0])
        optim_m = kwargs['optim']
        # scheduler_m = kwargs['scheduler']      
        optim_m.zero_grad()
        loss = self.criterion(preds, kwargs['batch'][1])
        loss.backward()
        optim_m.step()
        # scheduler_m.step()
        return loss.item();
  
      def evaluation_step(self,**kwargs):
        test_iters = kwargs['test_iters']
        target_scale = kwargs['target_scale']
        test_preds_sum = 0.
        for i in range(test_iters):
          self.train()
          test_preds_sum = test_preds_sum + (self(kwargs['test_tuple'][0])*target_scale)
        test_preds = test_preds_sum/test_iters
        test_score = kwargs['eval_func'](test_preds, kwargs['test_tuple'][1]*target_scale)
        scores = []
        for i in test_score:
          scores.append(i.detach().cpu().numpy())
  
        return scores