hunk.code_window: @register_model
   def vit_relpos_base_patch32_plus_rpn_256(pretrained=False, **kwargs):
       """ ViT-Base (ViT-B/32+) w/ relative log-coord position  and residual post-norm, no class token
       """
       model_kwargs = dict(
  
hunk.labels: keep keep keep keep replace keep keep
hunk.after_edit:     """ ViT-Base (ViT-B/32+) w/ relative log-coord position and residual post-norm, no class token
  
hunk.file_path: timm/models/vision_transformer_relpos.py
hunk.type: replace
file: """ Optimzier Tests
  
  These tests were adapted from PyTorch' optimizer tests.
  
  """
  import math
  import pytest
  import functools
  from copy import deepcopy
  
  import torch
  from torch.testing._internal.common_utils import TestCase
  from torch.autograd import Variable
  from timm.scheduler import PlateauLRScheduler
  
  from timm.optim import create_optimizer_v2
  
  
  # HACK relying on internal PyTorch test functionality for comparisons that I don't want to write
  torch_tc = TestCase()
  
  
  def _test_basic_cases_template(weight, bias, input, constructor, scheduler_constructors):
      weight = Variable(weight, requires_grad=True)
      bias = Variable(bias, requires_grad=True)
      input = Variable(input)
      optimizer = constructor(weight, bias)
      schedulers = []
      for scheduler_constructor in scheduler_constructors:
          schedulers.append(scheduler_constructor(optimizer))
  
      # to check if the optimizer can be printed as a string
      optimizer.__repr__()
  
      def fn():
          optimizer.zero_grad()
          y = weight.mv(input)
          if y.is_cuda and bias.is_cuda and y.get_device() != bias.get_device():
              y = y.cuda(bias.get_device())
          loss = (y + bias).pow(2).sum()
          loss.backward()
          return loss
  
      initial_value = fn().item()
      for _i in range(200):
          for scheduler in schedulers:
              if isinstance(scheduler, PlateauLRScheduler):
                  val_loss = fn()
                  scheduler.step(val_loss)
              else:
                  scheduler.step()
          optimizer.step(fn)
  
      assert fn().item() < initial_value
  
  
  def _test_state_dict(weight, bias, input, constructor):
      weight = Variable(weight, requires_grad=True)
      bias = Variable(bias, requires_grad=True)
      input = Variable(input)
  
      def fn_base(optimizer, weight, bias):
          optimizer.zero_grad()
          i = input_cuda if weight.is_cuda else input
          loss = (weight.mv(i) + bias).pow(2).sum()
          loss.backward()
          return loss
  
      optimizer = constructor(weight, bias)
      fn = functools.partial(fn_base, optimizer, weight, bias)
  
      # Prime the optimizer
      for _i in range(20):
          optimizer.step(fn)
      # Clone the weights and construct new optimizer for them
      weight_c = Variable(weight.data.clone(), requires_grad=True)
      bias_c = Variable(bias.data.clone(), requires_grad=True)
      optimizer_c = constructor(weight_c, bias_c)
      fn_c = functools.partial(fn_base, optimizer_c, weight_c, bias_c)
      # Load state dict
      state_dict = deepcopy(optimizer.state_dict())
      state_dict_c = deepcopy(optimizer.state_dict())
      optimizer_c.load_state_dict(state_dict_c)
  
      # Run both optimizations in parallel
      for _i in range(20):
          optimizer.step(fn)
          optimizer_c.step(fn_c)
          #assert torch.equal(weight, weight_c)
          #assert torch.equal(bias, bias_c)
          torch_tc.assertEqual(weight, weight_c)
          torch_tc.assertEqual(bias, bias_c)
      # Make sure state dict wasn't modified
      torch_tc.assertEqual(state_dict, state_dict_c)
      # Make sure state dict is deterministic with equal but not identical parameters
      torch_tc.assertEqual(optimizer.state_dict(), optimizer_c.state_dict())
      # Make sure repeated parameters have identical representation in state dict
      optimizer_c.param_groups.extend(optimizer_c.param_groups)
      torch_tc.assertEqual(optimizer.state_dict()['param_groups'][-1], optimizer_c.state_dict()['param_groups'][-1])
  
      # Check that state dict can be loaded even when we cast parameters
      # to a different type and move to a different device.
      if not torch.cuda.is_available():
          return
  
      input_cuda = Variable(input.data.float().cuda())
      weight_cuda = Variable(weight.data.float().cuda(), requires_grad=True)
      bias_cuda = Variable(bias.data.float().cuda(), requires_grad=True)
      optimizer_cuda = constructor(weight_cuda, bias_cuda)
      fn_cuda = functools.partial(fn_base, optimizer_cuda, weight_cuda, bias_cuda)
  
      state_dict = deepcopy(optimizer.state_dict())
      state_dict_c = deepcopy(optimizer.state_dict())
      optimizer_cuda.load_state_dict(state_dict_c)
  
      # Make sure state dict wasn't modified
      torch_tc.assertEqual(state_dict, state_dict_c)
  
      for _i in range(20):
          optimizer.step(fn)
          optimizer_cuda.step(fn_cuda)
          torch_tc.assertEqual(weight, weight_cuda)
          torch_tc.assertEqual(bias, bias_cuda)
  
      # validate deepcopy() copies all public attributes
      def getPublicAttr(obj):
          return set(k for k in obj.__dict__ if not k.startswith('_'))
  
      assert getPublicAttr(optimizer) == getPublicAttr(deepcopy(optimizer))
  
  
  def _test_basic_cases(constructor, scheduler_constructors=None):
      if scheduler_constructors is None:
          scheduler_constructors = []
      _test_state_dict(
          torch.randn(10, 5),
          torch.randn(10),
          torch.randn(5),
          constructor
      )
      _test_basic_cases_template(
          torch.randn(10, 5),
          torch.randn(10),
          torch.randn(5),
          constructor,
          scheduler_constructors
      )
      # non-contiguous parameters
      _test_basic_cases_template(
          torch.randn(10, 5, 2)[..., 0],
          torch.randn(10, 2)[..., 0],
          torch.randn(5),
          constructor,
          scheduler_constructors
      )
      # CUDA
      if not torch.cuda.is_available():
          return
      _test_basic_cases_template(
          torch.randn(10, 5).cuda(),
          torch.randn(10).cuda(),
          torch.randn(5).cuda(),
          constructor,
          scheduler_constructors
      )
  
  
  def _test_model(optimizer, params, device=torch.device('cpu')):
      weight = torch.tensor(
          [[-0.2109, -0.4976], [-0.1413, -0.3420], [-0.2524, 0.6976]],
          device=device, requires_grad=True)
      bias = torch.tensor([-0.1085, -0.2979, 0.6892], device=device, requires_grad=True)
      weight2 = torch.tensor([[-0.0508, -0.3941, -0.2843]], device=device, requires_grad=True)
      bias2 = torch.tensor([-0.0711], device=device, requires_grad=True)
      input = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6], device=device).reshape(3, 2)
  
      model = torch.nn.Sequential(torch.nn.Linear(2, 3),
                                  torch.nn.Sigmoid(),
                                  torch.nn.Linear(3, 1),
                                  torch.nn.Sigmoid())
      model.to(device)
  
      pretrained_dict = model.state_dict()
      pretrained_dict['0.weight'] = weight
      pretrained_dict['0.bias'] = bias
      pretrained_dict['2.weight'] = weight2
      pretrained_dict['2.bias'] = bias2
      model.load_state_dict(pretrained_dict)
  
      optimizer = create_optimizer_v2(model, opt=optimizer, **params)
  
      prev_loss = float('inf')
      for i in range(20):
          optimizer.zero_grad()
          output = model(input)
          loss = output.sum()
          loss.backward()
          loss = loss.item()
          assert loss < prev_loss
          prev_loss = loss
          optimizer.step()
  
  
  def rosenbrock(tensor):
      x, y = tensor
      return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2
  
  
  def drosenbrock(tensor):
      x, y = tensor
      return torch.tensor((-400 * x * (y - x ** 2) - 2 * (1 - x), 200 * (y - x ** 2)))
  
  
  def _test_rosenbrock(constructor, scheduler_constructors=None):
      if scheduler_constructors is None:
          scheduler_constructors = []
      params_t = torch.tensor([1.5, 1.5])
  
      params = Variable(params_t, requires_grad=True)
      optimizer = constructor([params])
      schedulers = []
      for scheduler_constructor in scheduler_constructors:
          schedulers.append(scheduler_constructor(optimizer))
  
      solution = torch.tensor([1, 1])
      initial_dist = params.data.dist(solution)
  
      def eval(params, w):
          # Depending on w, provide only the x or y gradient
          optimizer.zero_grad()
          loss = rosenbrock(params)
          loss.backward()
          grad = drosenbrock(params.data)
          # NB: We torture test the optimizer by returning an
          # uncoalesced sparse tensor
          if w:
              i = torch.LongTensor([[0, 0]])
              x = grad[0]
              v = torch.tensor([x / 4., x - x / 4.])
          else:
              i = torch.LongTensor([[1, 1]])
              y = grad[1]
              v = torch.tensor([y - y / 4., y / 4.])
          x = torch.sparse.DoubleTensor(i, v, torch.Size([2])).to(dtype=v.dtype)
          with torch.no_grad():
              params.grad = x.to_dense()
          return loss
  
      for i in range(2000):
          # Do cyclic coordinate descent
          w = i % 2
          optimizer.step(functools.partial(eval, params, w))
          for scheduler in schedulers:
              if isinstance(scheduler, PlateauLRScheduler):
                  scheduler.step(rosenbrock(params))
              else:
                  scheduler.step()
  
      torch_tc.assertLessEqual(params.data.dist(solution), initial_dist)
  
  
  def _build_params_dict(weight, bias, **kwargs):
      return [{'params': [weight]}, dict(params=[bias], **kwargs)]
  
  
  def _build_params_dict_single(weight, bias, **kwargs):
      return [dict(params=bias, **kwargs)]
  
  
  #@pytest.mark.parametrize('optimizer', ['sgd', 'momentum'])
  # FIXME momentum variant frequently fails in GitHub runner, but never local after many attempts
  @pytest.mark.parametrize('optimizer', ['sgd'])
  def test_sgd(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=1e-2),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=1e-2),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=1e-2), optimizer)
      )
      # _test_basic_cases(
      #     lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3),
      #     [lambda opt: StepLR(opt, gamma=0.9, step_size=10)]
      # )
      # _test_basic_cases(
      #     lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3),
      #     [lambda opt: WarmUpLR(opt, warmup_factor=0.4, warmup_iters=4, warmup_method="linear")]
      # )
      # _test_basic_cases(
      #     lambda weight, bias: optimizer([weight, bias], lr=1e-3),
      #     [lambda opt: WarmUpLR(opt, warmup_factor=0.4, warmup_iters=4, warmup_method="constant")]
      # )
      # _test_basic_cases(
      #     lambda weight, bias: optimizer([weight, bias], lr=1e-3),
      #     [lambda opt: StepLR(opt, gamma=0.9, step_size=10),
      #      lambda opt: WarmUpLR(opt, warmup_factor=0.4, warmup_iters=4)]
      # )
      # _test_basic_cases(
      #     lambda weight, bias: optimizer([weight, bias], lr=1e-3),
      #     [lambda opt: StepLR(opt, gamma=0.9, step_size=10),
      #      lambda opt: ReduceLROnPlateau(opt)]
      # )
      # _test_basic_cases(
      #     lambda weight, bias: optimizer([weight, bias], lr=1e-3),
      #     [lambda opt: StepLR(opt, gamma=0.99, step_size=10),
      #      lambda opt: ExponentialLR(opt, gamma=0.99),
      #      lambda opt: ReduceLROnPlateau(opt)]
      # )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=3e-3, momentum=1)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=3e-3, momentum=1, weight_decay=.1)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
      )
      _test_model(optimizer, dict(lr=1e-3))
  
  
  @pytest.mark.parametrize('optimizer',  ['adamw', 'adam', 'nadam', 'adamax'])
  def test_adam(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=5e-2)
      )
      _test_model(optimizer, dict(lr=5e-2))
  
  
  @pytest.mark.parametrize('optimizer',  ['adabelief'])
  def test_adabelief(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3, weight_decay=1)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=5e-2)
      )
      _test_model(optimizer, dict(lr=5e-2))
  
  
  @pytest.mark.parametrize('optimizer',  ['radam', 'radabelief'])
  def test_rectified(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
      )
      _test_model(optimizer, dict(lr=1e-3))
  
  
  @pytest.mark.parametrize('optimizer',   ['adadelta', 'adagrad'])
  def test_adaother(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3, weight_decay=1)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-1)
      )
      _test_model(optimizer, dict(lr=5e-2))
  
  
  @pytest.mark.parametrize('optimizer',   ['adafactor'])
  def test_adafactor(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(_build_params_dict_single(weight, bias), optimizer)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3, weight_decay=1)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=5e-2)
      )
      _test_model(optimizer, dict(lr=5e-2))
  
  
  @pytest.mark.parametrize('optimizer',  ['lamb', 'lambc'])
  def test_lamb(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=1e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=1e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=1e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
      )
      _test_model(optimizer, dict(lr=1e-3))
  
  
  @pytest.mark.parametrize('optimizer',  ['lars', 'larc', 'nlars', 'nlarc'])
  def test_lars(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=1e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=1e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=1e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
      )
      _test_model(optimizer, dict(lr=1e-3))
  
  
  @pytest.mark.parametrize('optimizer',  ['madgrad', 'madgradw'])
  def test_madgrad(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-2)
      )
      _test_model(optimizer, dict(lr=1e-2))
  
  
  @pytest.mark.parametrize('optimizer',  ['novograd'])
  def test_novograd(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
      )
      _test_model(optimizer, dict(lr=1e-3))
  
  
  @pytest.mark.parametrize('optimizer', ['rmsprop', 'rmsproptf'])
  def test_rmsprop(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-2)
      )
      _test_model(optimizer, dict(lr=1e-2))
  
  
  @pytest.mark.parametrize('optimizer', ['adamp'])
  def test_adamp(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=5e-2)
      )
      _test_model(optimizer, dict(lr=5e-2))
  
  
  @pytest.mark.parametrize('optimizer', ['sgdp'])
  def test_sgdp(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
      )
      _test_model(optimizer, dict(lr=1e-3))
  
  
  @pytest.mark.parametrize('optimizer', ['lookahead_sgd', 'lookahead_momentum'])
  def test_lookahead_sgd(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-3)
      )
  
  
  @pytest.mark.parametrize('optimizer', ['lookahead_adamw', 'lookahead_adam'])
  def test_lookahead_adam(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=5e-2)
      )
  
  
  @pytest.mark.parametrize('optimizer', ['lookahead_radam'])
  def test_lookahead_radam(optimizer):
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2([weight, bias], optimizer, lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3),
              optimizer,
              lr=1e-3)
      )
      _test_basic_cases(
          lambda weight, bias: create_optimizer_v2(
              _build_params_dict_single(weight, bias, lr=3e-3), optimizer)
      )
      _test_rosenbrock(
          lambda params: create_optimizer_v2(params, optimizer, lr=1e-4)
      )
  
  