jettify · jstjohn · Oct 27, 2021 · Oct 28, 2021 · Oct 28, 2021 · Nov 2, 2021
diff --git a/README.rst b/README.rst
diff --git a/docs/index.rst b/docs/index.rst
@@ -60,6 +60,9 @@ Supported Optimizers
 | :ref:`Adafactor`| https://arxiv.org/abs/1804.04235                                              |
 +-----------------+-------------------------------------------------------------------------------+
 |                 |                                                                               |
+| :ref:`AdamD`    | https://arxiv.org/abs/2110.10828                                             |
++-----------------+-------------------------------------------------------------------------------+
+|                 |                                                                               |
 | :ref:`AdamP`    | https://arxiv.org/abs/1804.00325                                              |
 +-----------------+-------------------------------------------------------------------------------+
 |                 |                                                                               |

diff --git a/docs/rastrigin_A2GradExp.png b/docs/rastrigin_A2GradExp.png
diff --git a/docs/rastrigin_A2GradInc.png b/docs/rastrigin_A2GradInc.png
diff --git a/docs/rastrigin_A2GradUni.png b/docs/rastrigin_A2GradUni.png
diff --git a/docs/rastrigin_AccSGD.png b/docs/rastrigin_AccSGD.png
diff --git a/docs/rastrigin_AdaBelief.png b/docs/rastrigin_AdaBelief.png
diff --git a/docs/rastrigin_AdaBelief_adamD.png b/docs/rastrigin_AdaBelief_adamD.png
diff --git a/docs/rastrigin_AdaBound.png b/docs/rastrigin_AdaBound.png
diff --git a/docs/rastrigin_AdaBound_adamD.png b/docs/rastrigin_AdaBound_adamD.png
diff --git a/docs/rastrigin_AdaMod.png b/docs/rastrigin_AdaMod.png
diff --git a/docs/rastrigin_AdaMod_adamD.png b/docs/rastrigin_AdaMod_adamD.png
diff --git a/docs/rastrigin_Adafactor.png b/docs/rastrigin_Adafactor.png
diff --git a/docs/rastrigin_Adahessian.png b/docs/rastrigin_Adahessian.png
diff --git a/docs/rastrigin_Adahessian_adamD.png b/docs/rastrigin_Adahessian_adamD.png
diff --git a/docs/rastrigin_Adam.png b/docs/rastrigin_Adam.png
diff --git a/docs/rastrigin_AdamP.png b/docs/rastrigin_AdamP.png
diff --git a/docs/rastrigin_AdamP_adamD.png b/docs/rastrigin_AdamP_adamD.png
diff --git a/docs/rastrigin_AdamW_internal.png b/docs/rastrigin_AdamW_internal.png
diff --git a/docs/rastrigin_AdamW_internal_adamD.png b/docs/rastrigin_AdamW_internal_adamD.png
diff --git a/docs/rastrigin_Adam_internal.png b/docs/rastrigin_Adam_internal.png
diff --git a/docs/rastrigin_Adam_internal_adamD.png b/docs/rastrigin_Adam_internal_adamD.png
diff --git a/docs/rastrigin_AggMo.png b/docs/rastrigin_AggMo.png
diff --git a/docs/rastrigin_Apollo.png b/docs/rastrigin_Apollo.png
diff --git a/docs/rastrigin_DiffGrad.png b/docs/rastrigin_DiffGrad.png
diff --git a/docs/rastrigin_DiffGrad_adamD.png b/docs/rastrigin_DiffGrad_adamD.png
diff --git a/docs/rastrigin_Lamb.png b/docs/rastrigin_Lamb.png
diff --git a/docs/rastrigin_Lamb_adamD.png b/docs/rastrigin_Lamb_adamD.png
diff --git a/docs/rastrigin_LookaheadYogi.png b/docs/rastrigin_LookaheadYogi.png
diff --git a/docs/rastrigin_MADGRAD.png b/docs/rastrigin_MADGRAD.png
diff --git a/docs/rastrigin_NovoGrad.png b/docs/rastrigin_NovoGrad.png
diff --git a/docs/rastrigin_PID.png b/docs/rastrigin_PID.png
diff --git a/docs/rastrigin_QHAdam.png b/docs/rastrigin_QHAdam.png
diff --git a/docs/rastrigin_QHM.png b/docs/rastrigin_QHM.png
diff --git a/docs/rastrigin_Ranger.png b/docs/rastrigin_Ranger.png
diff --git a/docs/rastrigin_RangerQH.png b/docs/rastrigin_RangerQH.png
diff --git a/docs/rastrigin_RangerVA.png b/docs/rastrigin_RangerVA.png
diff --git a/docs/rastrigin_SGD.png b/docs/rastrigin_SGD.png
diff --git a/docs/rastrigin_SGDP.png b/docs/rastrigin_SGDP.png
diff --git a/docs/rastrigin_SGDW.png b/docs/rastrigin_SGDW.png
diff --git a/docs/rastrigin_SWATS.png b/docs/rastrigin_SWATS.png
diff --git a/docs/rastrigin_SWATS_adamD.png b/docs/rastrigin_SWATS_adamD.png
diff --git a/docs/rastrigin_Shampoo.png b/docs/rastrigin_Shampoo.png
diff --git a/docs/rastrigin_Yogi.png b/docs/rastrigin_Yogi.png
diff --git a/docs/rastrigin_Yogi_adamD.png b/docs/rastrigin_Yogi_adamD.png
diff --git a/docs/rosenbrock_A2GradExp.png b/docs/rosenbrock_A2GradExp.png
diff --git a/docs/rosenbrock_A2GradInc.png b/docs/rosenbrock_A2GradInc.png
diff --git a/docs/rosenbrock_A2GradUni.png b/docs/rosenbrock_A2GradUni.png
diff --git a/docs/rosenbrock_AccSGD.png b/docs/rosenbrock_AccSGD.png
diff --git a/docs/rosenbrock_AdaBelief.png b/docs/rosenbrock_AdaBelief.png
diff --git a/docs/rosenbrock_AdaBelief_adamD.png b/docs/rosenbrock_AdaBelief_adamD.png
diff --git a/docs/rosenbrock_AdaBound.png b/docs/rosenbrock_AdaBound.png
diff --git a/docs/rosenbrock_AdaBound_adamD.png b/docs/rosenbrock_AdaBound_adamD.png
diff --git a/docs/rosenbrock_AdaMod.png b/docs/rosenbrock_AdaMod.png
diff --git a/docs/rosenbrock_AdaMod_adamD.png b/docs/rosenbrock_AdaMod_adamD.png
diff --git a/docs/rosenbrock_Adafactor.png b/docs/rosenbrock_Adafactor.png
diff --git a/docs/rosenbrock_Adahessian.png b/docs/rosenbrock_Adahessian.png
diff --git a/docs/rosenbrock_Adahessian_adamD.png b/docs/rosenbrock_Adahessian_adamD.png
diff --git a/docs/rosenbrock_Adam.png b/docs/rosenbrock_Adam.png
diff --git a/docs/rosenbrock_AdamP.png b/docs/rosenbrock_AdamP.png
diff --git a/docs/rosenbrock_AdamP_adamD.png b/docs/rosenbrock_AdamP_adamD.png
diff --git a/docs/rosenbrock_AdamW_internal.png b/docs/rosenbrock_AdamW_internal.png
diff --git a/docs/rosenbrock_AdamW_internal_adamD.png b/docs/rosenbrock_AdamW_internal_adamD.png
diff --git a/docs/rosenbrock_Adam_internal.png b/docs/rosenbrock_Adam_internal.png
diff --git a/docs/rosenbrock_Adam_internal_adamD.png b/docs/rosenbrock_Adam_internal_adamD.png
diff --git a/docs/rosenbrock_AggMo.png b/docs/rosenbrock_AggMo.png
diff --git a/docs/rosenbrock_Apollo.png b/docs/rosenbrock_Apollo.png
diff --git a/docs/rosenbrock_DiffGrad.png b/docs/rosenbrock_DiffGrad.png
diff --git a/docs/rosenbrock_DiffGrad_adamD.png b/docs/rosenbrock_DiffGrad_adamD.png
diff --git a/docs/rosenbrock_Lamb.png b/docs/rosenbrock_Lamb.png
diff --git a/docs/rosenbrock_Lamb_adamD.png b/docs/rosenbrock_Lamb_adamD.png
diff --git a/docs/rosenbrock_LookaheadYogi.png b/docs/rosenbrock_LookaheadYogi.png
diff --git a/docs/rosenbrock_MADGRAD.png b/docs/rosenbrock_MADGRAD.png
diff --git a/docs/rosenbrock_NovoGrad.png b/docs/rosenbrock_NovoGrad.png
diff --git a/docs/rosenbrock_PID.png b/docs/rosenbrock_PID.png
diff --git a/docs/rosenbrock_QHAdam.png b/docs/rosenbrock_QHAdam.png
diff --git a/docs/rosenbrock_QHM.png b/docs/rosenbrock_QHM.png
diff --git a/docs/rosenbrock_Ranger.png b/docs/rosenbrock_Ranger.png
diff --git a/docs/rosenbrock_RangerQH.png b/docs/rosenbrock_RangerQH.png
diff --git a/docs/rosenbrock_RangerVA.png b/docs/rosenbrock_RangerVA.png
diff --git a/docs/rosenbrock_SGD.png b/docs/rosenbrock_SGD.png
diff --git a/docs/rosenbrock_SGDP.png b/docs/rosenbrock_SGDP.png
diff --git a/docs/rosenbrock_SGDW.png b/docs/rosenbrock_SGDW.png
diff --git a/docs/rosenbrock_SWATS.png b/docs/rosenbrock_SWATS.png
diff --git a/docs/rosenbrock_SWATS_adamD.png b/docs/rosenbrock_SWATS_adamD.png
diff --git a/docs/rosenbrock_Shampoo.png b/docs/rosenbrock_Shampoo.png
diff --git a/docs/rosenbrock_Yogi.png b/docs/rosenbrock_Yogi.png
diff --git a/docs/rosenbrock_Yogi_adamD.png b/docs/rosenbrock_Yogi_adamD.png
diff --git a/examples/viz_optimizers.py b/examples/viz_optimizers.py
@@ -1,4 +1,5 @@
 import math
+from typing import List, Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -9,6 +10,9 @@
 
 plt.style.use('seaborn-white')
 
+NUM_ITER: int = 500
+NUM_ITER_HPARAM: int = 200
+
 
 def rosenbrock(tensor):
     # https://en.wikipedia.org/wiki/Test_functions_for_optimization
@@ -29,11 +33,10 @@ def rastrigin(tensor, lib=torch):
 
 
 def execute_steps(
-    func, initial_state, optimizer_class, optimizer_config, num_iter=500
+    func, initial_state, optimizer_class, optimizer_config, num_iter=NUM_ITER
 ):
     x = torch.Tensor(initial_state).requires_grad_(True)
     optimizer = optimizer_class([x], **optimizer_config)
-    steps = []
     steps = np.zeros((2, num_iter + 1))
     steps[:, 0] = np.array(initial_state)
     for i in range(1, num_iter + 1):
@@ -49,10 +52,11 @@ def execute_steps(
 def objective_rastrigin(params):
     lr = params['lr']
     optimizer_class = params['optimizer_class']
+    kwargs = params['kwargs']
     initial_state = (-2.0, 3.5)
     minimum = (0, 0)
-    optimizer_config = dict(lr=lr)
-    num_iter = 100
+    optimizer_config = dict(lr=lr, **kwargs)
+    num_iter = NUM_ITER_HPARAM
     steps = execute_steps(
         rastrigin, initial_state, optimizer_class, optimizer_config, num_iter
     )
@@ -62,74 +66,167 @@ def objective_rastrigin(params):
 def objective_rosenbrok(params):
     lr = params['lr']
     optimizer_class = params['optimizer_class']
+    kwargs = params['kwargs']
     minimum = (1.0, 1.0)
     initial_state = (-2.0, 2.0)
-    optimizer_config = dict(lr=lr)
-    num_iter = 100
+    optimizer_config = dict(lr=lr, **kwargs)
+    num_iter = NUM_ITER_HPARAM
     steps = execute_steps(
         rosenbrock, initial_state, optimizer_class, optimizer_config, num_iter
     )
     return (steps[0][-1] - minimum[0]) ** 2 + (steps[1][-1] - minimum[1]) ** 2
 
 
-def plot_rastrigin(grad_iter, optimizer_name, lr):
+def plot_rastrigin(grad_iters, optimizer_name, lr):
     x = np.linspace(-4.5, 4.5, 250)
     y = np.linspace(-4.5, 4.5, 250)
     minimum = (0, 0)
 
     X, Y = np.meshgrid(x, y)
     Z = rastrigin([X, Y], lib=np)
-
-    iter_x, iter_y = grad_iter[0, :], grad_iter[1, :]
-
+    assert len(grad_iters) <= 3, "Cannot handle more than three states"
+    l = None
     fig = plt.figure(figsize=(8, 8))
 
     ax = fig.add_subplot(1, 1, 1)
-    ax.contour(X, Y, Z, 20, cmap='jet')
-    ax.plot(iter_x, iter_y, color='r', marker='x')
-    ax.set_title(
+    plt.contour(X, Y, Z, 20, cmap='jet', alpha=0.75)
+    for grad_iter, color in zip(grad_iters, ['r', 'm', 'c']):
+        iter_x, iter_y = grad_iter[0, :], grad_iter[1, :]
+        if l is None:
+            l = len(iter_x)
+        ax.plot(iter_x, iter_y, color=color, marker=None, alpha=0.75)
+        for px, py, pdx, pdy in zip(
+            iter_x[:-1],
+            iter_y[:-1],
+            iter_x[1:] - iter_x[:-1],
+            iter_y[1:] - iter_y[:-1],
+        ):
+            ax.arrow(
+                x=px,
+                y=py,
+                dx=pdx,
+                dy=pdy,
+                overhang=0.5,
+                width=0.001,
+                head_width=0.08,
+                length_includes_head=True,
+                color=color,
+                visible=True,
+            )
+        # Starting point
+        ax.plot(
+            iter_x[0],
+            iter_y[0],
+            marker="s",
+            markersize=11,
+            markeredgecolor="black",
+            markerfacecolor=color,
+            markeredgewidth=2,
+        )
+        # Ending point
+        ax.plot(
+            iter_x[-1],
+            iter_y[-1],
+            marker="P",
+            markersize=11,
+            markeredgecolor="black",
+            markerfacecolor=color,
+            markeredgewidth=2,
+        )
+    plt.title(
         'Rastrigin func: {} with '
-        '{} iterations, lr={:.6}'.format(optimizer_name, len(iter_x), lr)
+        '{} iterations, lr={:.6}'.format(optimizer_name, l, lr)
     )
-    plt.plot(*minimum, 'gD')
-    plt.plot(iter_x[-1], iter_y[-1], 'rD')
+    plt.xlim(-4.5, 4.5)
+    plt.ylim(-4.5, 4.5)
+    plt.plot(*minimum, 'X', color="green", markersize=11)
     plt.savefig('docs/rastrigin_{}.png'.format(optimizer_name))
+    plt.close()
 
 
-def plot_rosenbrok(grad_iter, optimizer_name, lr):
+def plot_rosenbrok(grad_iters, optimizer_name, lr):
     x = np.linspace(-2, 2, 250)
     y = np.linspace(-1, 3, 250)
     minimum = (1.0, 1.0)
+    assert len(grad_iters) <= 3, "Cannot handle more than three states"
 
     X, Y = np.meshgrid(x, y)
     Z = rosenbrock([X, Y])
-
-    iter_x, iter_y = grad_iter[0, :], grad_iter[1, :]
-
+    l = None
     fig = plt.figure(figsize=(8, 8))
-
     ax = fig.add_subplot(1, 1, 1)
-    ax.contour(X, Y, Z, 90, cmap='jet')
-    ax.plot(iter_x, iter_y, color='r', marker='x')
-
-    ax.set_title(
+    ax.contour(X, Y, Z, 90, cmap='jet', alpha=0.75)
+    for grad_iter, color in zip(grad_iters, ['r', 'm', 'c']):
+        iter_x, iter_y = grad_iter[0, :], grad_iter[1, :]
+        if l is None:
+            l = len(iter_x)
+        ax.plot(iter_x, iter_y, color=color, marker=None, alpha=0.75)
+
+        for px, py, pdx, pdy in zip(
+            iter_x[:-1],
+            iter_y[:-1],
+            iter_x[1:] - iter_x[:-1],
+            iter_y[1:] - iter_y[:-1],
+        ):
+            ax.arrow(
+                x=px,
+                y=py,
+                dx=pdx,
+                dy=pdy,
+                overhang=0.5,
+                width=0.001,
+                head_width=0.0375,
+                length_includes_head=True,
+                color=color,
+                visible=True,
+            )
+        # Starting point
+        ax.plot(
+            iter_x[0],
+            iter_y[0],
+            marker="s",
+            markersize=11,
+            markeredgecolor="black",
+            markerfacecolor=color,
+            markeredgewidth=2,
+        )
+        # Ending point
+        ax.plot(
+            iter_x[-1],
+            iter_y[-1],
+            marker="P",
+            markersize=11,
+            markeredgecolor="black",
+            markerfacecolor=color,
+            markeredgewidth=2,
+        )
+    plt.title(
         'Rosenbrock func: {} with {} '
-        'iterations, lr={:.6}'.format(optimizer_name, len(iter_x), lr)
+        'iterations, lr={:.6}'.format(optimizer_name, l, lr)
     )
-    plt.plot(*minimum, 'gD')
-    plt.plot(iter_x[-1], iter_y[-1], 'rD')
+    plt.plot(*minimum, 'X', color="green", markersize=11)
+    plt.xlim(-2, 2)
+    plt.ylim(-1, 3)
     plt.savefig('docs/rosenbrock_{}.png'.format(optimizer_name))
+    plt.close()
 
 
 def execute_experiments(
-    optimizers, objective, func, plot_func, initial_state, seed=1
+    optimizers,
+    objective,
+    func,
+    plot_func,
+    initial_states: List[Tuple[float, float]],
+    seed=1,
 ):
     seed = seed
     for item in optimizers:
-        optimizer_class, lr_low, lr_hi = item
+        optimizer_class, lr_low, lr_hi, kwargs, extra_desc = item
+        extra_desc_str = '' if not extra_desc else f'_{extra_desc}'
         space = {
             'optimizer_class': hp.choice('optimizer_class', [optimizer_class]),
             'lr': hp.loguniform('lr', lr_low, lr_hi),
+            'kwargs': kwargs,
         }
         best = fmin(
             fn=objective,
@@ -139,15 +236,21 @@ def execute_experiments(
             rstate=np.random.RandomState(seed),
         )
         print(best['lr'], optimizer_class)
-
-        steps = execute_steps(
-            func,
-            initial_state,
-            optimizer_class,
-            {'lr': best['lr']},
-            num_iter=500,
+        steps_lst = []
+        for initial_state in initial_states:
+            steps = execute_steps(
+                func,
+                initial_state,
+                optimizer_class,
+                {'lr': best['lr'], **kwargs},
+                num_iter=NUM_ITER,
+            )
+            steps_lst.append(steps)
+        plot_func(
+            steps_lst,
+            f'{optimizer_class.__name__}{extra_desc_str}',
+            best['lr'],
         )
-        plot_func(steps, optimizer_class.__name__, best['lr'])
 
 
 def LookaheadYogi(*a, **kw):
@@ -162,52 +265,84 @@ def LookaheadYogi(*a, **kw):
     # help to converge on better lr faster.
     optimizers = [
         # baselines
-        (torch.optim.Adam, -8, 0.5),
-        (torch.optim.SGD, -8, -1.0),
+        (torch.optim.Adam, -8, 0.5, {}, None),
+        (torch.optim.SGD, -8, -1.0, {}, None),
         # Adam based
-        (optim.AdaBound, -8, 0.3),
-        (optim.Adahessian, -1, 8),
-        (optim.AdaMod, -8, 0.2),
-        (optim.AdamP, -8, 0.2),
-        (optim.DiffGrad, -8, 0.4),
-        (optim.Lamb, -8, -2.9),
-        (optim.MADGRAD, -8, 0.5),
-        (optim.NovoGrad, -8, -1.7),
-        (optim.RAdam, -8, 0.5),
-        (optim.Yogi, -8, 0.1),
+        (optim.Adam, -8, 0.5, {}, 'internal'),
+        (
+            optim.Adam,
+            -8,
+            0.5,
+            {'adamd_bias_correction': True},
+            'internal_adamD',
+        ),
+        (optim.AdamW, -8, 0.5, {}, 'internal'),
+        (
+            optim.AdamW,
+            -8,
+            0.5,
+            {'adamd_bias_correction': True},
+            'internal_adamD',
+        ),
+        (optim.AdaBound, -8, 0.3, {}, None),
+        (optim.AdaBound, -8, 0.3, {'adamd_bias_correction': True}, 'adamD'),
+        # TODO
+        (optim.Adahessian, -1, 8, {}, None),
+        (optim.Adahessian, -1, 8, {'adamd_bias_correction': True}, 'adamD'),
+        (optim.AdaMod, -8, 0.2, {}, None),
+        (optim.AdaMod, -8, 0.2, {'adamd_bias_correction': True}, 'adamD'),
+        (optim.AdamP, -8, 0.2, {}, None),
+        (optim.AdamP, -8, 0.2, {'adamd_bias_correction': True}, 'adamD'),
+        (optim.DiffGrad, -8, 0.4, {}, None),
+        (optim.DiffGrad, -8, 0.4, {'adamd_bias_correction': True}, 'adamD'),
+        (optim.Lamb, -8, -2.9, {}, None),
+        (
+            optim.Lamb,
+            -8,
+            -2.9,
+            {'debias': True, 'adamd_bias_correction': True},
+            'adamD',
+        ),
+        (optim.MADGRAD, -8, 0.5, {}, None),
+        (optim.NovoGrad, -8, -1.7, {}, None),
+        (optim.RAdam, -8, 0.5, {}, None),
+        (optim.Yogi, -8, 0.1, {}, None),
+        (optim.Yogi, -8, 0.1, {'adamd_bias_correction': True}, 'adamD'),
         # SGD/Momentum based
-        (optim.AccSGD, -8, -1.4),
-        (optim.SGDW, -8, -1.5),
-        (optim.SGDP, -8, -1.5),
-        (optim.PID, -8, -1.0),
-        (optim.QHM, -6, -0.2),
-        (optim.QHAdam, -8, 0.1),
-        (optim.Ranger, -8, 0.1),
-        (optim.RangerQH, -8, 0.1),
-        (optim.RangerVA, -8, 0.1),
-        (optim.Shampoo, -8, 0.1),
-        (LookaheadYogi, -8, 0.1),
-        (optim.AggMo, -8, -1.5),
-        (optim.SWATS, -8, -1.5),
-        (optim.Adafactor, -8, 0.5),
-        (optim.A2GradUni, -8, 0.1),
-        (optim.A2GradInc, -8, 0.1),
-        (optim.A2GradExp, -8, 0.1),
-        (optim.AdaBelief, -8, 0.1),
-        (optim.Apollo, -8, 0.1),
+        (optim.AccSGD, -8, -1.4, {}, None),
+        (optim.SGDW, -8, -1.5, {}, None),
+        (optim.SGDP, -8, -1.5, {}, None),
+        (optim.PID, -8, -1.0, {}, None),
+        (optim.QHM, -6, -0.2, {}, None),
+        (optim.QHAdam, -8, 0.1, {}, None),
+        (optim.Ranger, -8, 0.1, {}, None),
+        (optim.RangerQH, -8, 0.1, {}, None),
+        (optim.RangerVA, -8, 0.1, {}, None),
+        (optim.Shampoo, -8, 0.1, {}, None),
+        (LookaheadYogi, -8, 0.1, {}, None),
+        (optim.AggMo, -8, -1.5, {}, None),
+        (optim.SWATS, -8, -1.5, {}, None),
+        (optim.SWATS, -8, -1.5, {'adamd_bias_correction': True}, 'adamD'),
+        (optim.Adafactor, -8, 0.5, {}, None),
+        (optim.A2GradUni, -8, 0.1, {}, None),
+        (optim.A2GradInc, -8, 0.1, {}, None),
+        (optim.A2GradExp, -8, 0.1, {}, None),
+        (optim.AdaBelief, -8, 0.1, {}, None),
+        (optim.AdaBelief, -8, 0.1, {'adamd_bias_correction': True}, 'adamD'),
+        (optim.Apollo, -8, 0.1, {}, None),
     ]
     execute_experiments(
         optimizers,
         objective_rastrigin,
         rastrigin,
         plot_rastrigin,
-        (-2.0, 3.5),
+        [(-2.0, 3.5), (1.0, -2.0)],
     )
 
     execute_experiments(
         optimizers,
         objective_rosenbrok,
         rosenbrock,
         plot_rosenbrok,
-        (-2.0, 2.0),
+        [(-2.0, 2.0), (-0.5, 2.75)],
     )
diff --git a/tests/__init__.py b/tests/__init__.py
-Original file line number
+Diff line change
@@ Expand Up / @@ -60,6 +60,9 @@ Supported Optimizers @@
     | :ref:`Adafactor`| https://arxiv.org/abs/1804.04235                                              |
     +-----------------+-------------------------------------------------------------------------------+
     |                 |                                                                               |
+    | :ref:`AdamD`    | https://arxiv.org/abs/2110.10828                                             |
+    +-----------------+-------------------------------------------------------------------------------+
+    |                 |                                                                               |
     | :ref:`AdamP`    | https://arxiv.org/abs/1804.00325                                              |
     +-----------------+-------------------------------------------------------------------------------+
     |                 |                                                                               |
@@ Expand Down @@