text: from torch.nn.utils import clip_grad_norm_
  
  from modules.modeling_auxiLearn import Hypergrad, MTNgrad
  import torch
  
  from torch.nn.utils import clip_grad_norm_
  
  from modules.modeling_MELTR import AID_FP, VRBO, AID_CG, STOCBIO, STOCBIO_def
  from modules.modeling_auxiLearn import Hypergrad, MTNgrad
  import torch
  
  class MELTROptimizer_AID:
      def __init__(self, meta_optimizer, hpo_lr, truncate_iter=3, max_grad_norm=10, meltr_type=0):
          self.meta_optimizer = meta_optimizer
          if meltr_type == 0:
              self.hypergrad = MTNgrad(learning_rate=hpo_lr, truncate_iter=truncate_iter)
          elif meltr_type == 1:
              self.hypergrad = AID_FP()
          elif meltr_type == 2:
              self.hypergrad = AID_CG()
          self.max_grad_norm = max_grad_norm
  
      def step(self, train_loss, val_loss, parameters, aux_params):
          self.zero_grad()
  
          hyper_gards = self.hypergrad.grad(
              loss_val=val_loss,
              loss_train=train_loss,
              aux_params=aux_params,
              params=parameters,
          )
  
          for p, g in zip(aux_params, hyper_gards):
              if g is not None:
                  p.grad = -g
  
          if self.max_grad_norm is not None:
              clip_grad_norm_(aux_params, max_norm=self.max_grad_norm)
  
          self.meta_optimizer.step()
          # return sum([torch.norm(l) for l in hyper_gards]), sum([torch.norm(rl.grad) for rl in aux_params])
          return 0,0
      def zero_grad(self):
          self.meta_optimizer.zero_grad()
  
  
  class MELTROptimizer_2:
      def __init__(self, meta_optimizer, hpo_lr, truncate_iter=3, max_grad_norm=10, meltr_type=0):
          self.meta_optimizer = meta_optimizer
          self.meltr_type = meltr_type
          # if meltr_type == 3:
          #     self.hypergrad = STOCBIO()
          # elif meltr_type == 4:
          #     self.hypergrad = VRBO()
  
          self.max_grad_norm = max_grad_norm
  
      def step(self, args, net, auxiliary_combine_net, batchs):
          aux_params = auxiliary_combine_net.parameters()
  
          self.zero_grad()
  
          # hyper_gards = self.hypergrad.grad(
          #     args=args,
          #     net=net,
          #     auxiliary_combine_net=auxiliary_combine_net,
          #     batchs=batchs
          # )
          if self.meltr_type == 3:
              hyper_gards = STOCBIO_def(
                                  args=args,
                                  net=net,
                                  auxiliary_combine_net=auxiliary_combine_net,
                                  batchs=batchs)
  
  
          for p, g in zip(aux_params, hyper_gards):
              if g is not None:
                  p.grad = -g
  
          if self.max_grad_norm is not None:
              clip_grad_norm_(aux_params, max_norm=self.max_grad_norm)
  
          self.meta_optimizer.step()
          # return sum([torch.norm(l) for l in hyper_gards]), sum([torch.norm(rl.grad) for rl in aux_params])
          return 0, 0
  
      def zero_grad(self):
          self.meta_optimizer.zero_grad()
  
  
  class MELTROptimizer:
  
      def __init__(self, meta_optimizer, hpo_lr, truncate_iter=3, max_grad_norm=10, meltr_type=0):
          self.meta_optimizer = meta_optimizer
          if meltr_type == 0:
              self.hypergrad = MTNgrad(learning_rate=hpo_lr, truncate_iter=truncate_iter)
          elif meltr_type == 1:
              self.hypergrad = AID_FP()
  
          self.max_grad_norm = max_grad_norm
  
      def step(self, train_loss, val_loss, parameters, aux_params):
          self.zero_grad()
  
          hyper_gards = self.hypergrad.grad(
              loss_val=val_loss,
              loss_train=train_loss,
              aux_params=aux_params,
              params=parameters,
          )
          for p, g in zip(aux_params, hyper_gards):
              if g is not None:
                  p.grad = -g
  
          if self.max_grad_norm is not None:
              clip_grad_norm_(aux_params, max_norm=self.max_grad_norm)
  
          self.meta_optimizer.step()
          # return sum([torch.norm(l) for l in hyper_gards]), sum([torch.norm(rl.grad) for rl in aux_params])
          return 0,0
      def zero_grad(self):
          self.meta_optimizer.zero_grad()
  
  
  