code: def test_basic_with_learning_rate_decay():
      for i, dtype in enumerate(
          _dtypes_to_test(use_gpu=test_utils.is_gpu_available())
      ):
          # Initialize variables for numpy implementation.
          m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
          var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
          grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
          var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
          grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)
  
          var0 = tf.Variable(var0_np, name="var0_%d" % i)
          var1 = tf.Variable(var1_np, name="var1_%d" % i)
          grads0 = tf.constant(grads0_np)
          grads1 = tf.constant(grads1_np)
  
          learning_rate = 0.001
          beta_1 = 0.9
          beta_2 = 0.999
          epsilon = 1e-7
          decay = 0.5
          lamb_wd = 0.01
  
          opt = lamb.LAMB(
              learning_rate=learning_rate,
              beta_1=beta_1,
              beta_2=beta_2,
              epsilon=epsilon,
              weight_decay=lamb_wd,
              decay=decay,
          )
  
          # Run 3 steps of LAMB
          for t in range(3):
              opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
  
              lr_np = learning_rate / (1 + decay * t)
  
              var0_np, m0, v0 = lamb_update_numpy(
                  var0_np,
                  grads0_np,
                  t,
                  m0,
                  v0,
                  lr=lr_np,
                  lamb_wd=lamb_wd,
              )
              var1_np, m1, v1 = lamb_update_numpy(
                  var1_np,
                  grads1_np,
                  t,
                  m1,
                  v1,
                  lr=lr_np,
                  lamb_wd=lamb_wd,
              )
  
              # Validate updated params
              test_utils.assert_allclose_according_to_type(
                  var0_np, var0.numpy()
              )
              test_utils.assert_allclose_according_to_type(
                  var1_np, var1.numpy()
              )
  
imports: import numpy
  from numpy import linalg
  import pytest
  import tensorflow
  from tensorflow_addons.optimizers import lamb
  from tensorflow_addons.utils import test_utils
label_desc: Use image node_id 5 for calling a global function with example usage: test_basic_with_learning_rate_decay() without return types