code: def plot_rastrigin(grad_iter, 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, :]
  
      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(
          "Rastrigin func: {} with "
          "{} iterations, lr={:.6}".format(
              optimizer_name, len(iter_x), lr
          )
      )
      plt.plot(*minimum, "gD")
      plt.plot(iter_x[-1], iter_y[-1], "rD")
      plt.savefig("docs/rastrigin_{}.png".format(optimizer_name))
  
file_path: pytorch-optimizer/examples/viz_optimizers.py
imports: import math
  import matplotlib.pyplot
  import numpy
  import torch
  from hyperopt import fmin, hp, tpe
  import torch_optimizer
label_desc: Use image node_id 6 for calling a global function with example usage: plot_rastrigin(grad_iter, optimizer_name, lr) without return types