code: def test_basic(use_callable_params, dtype):
      # Initialize tf 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)
      var1 = tf.Variable(var1_np)
      grads0 = tf.constant(grads0_np)
      grads1 = tf.constant(grads1_np)
  
      def learning_rate():
          return 0.001
  
      if not use_callable_params:
          learning_rate = learning_rate()
  
      opt = lazy_adam.LazyAdam(learning_rate=learning_rate)
  
      # Run 3 steps of Adam
      for t in range(3):
          beta_1_power, beta_2_power = get_beta_accumulators(opt, dtype)
          test_utils.assert_allclose_according_to_type(
              0.9 ** (t + 1), beta_1_power
          )
          test_utils.assert_allclose_according_to_type(
              0.999 ** (t + 1), beta_2_power
          )
          opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
  
          var0_np, m0, v0 = adam_update_numpy(
              var0_np, grads0_np, t, m0, v0
          )
          var1_np, m1, v1 = adam_update_numpy(
              var1_np, grads1_np, t, m1, v1
          )
  
          # 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()
          )
  
file_path: addons/tensorflow_addons/optimizers/tests/lazy_adam_test.py
imports: import numpy
  import tensorflow
  from tensorflow_addons.optimizers import lazy_adam
  from tensorflow_addons.utils import test_utils
  import pytest
label_desc: Use image node_id 7 for calling a global function with example usage: test_basic(use_callable_params, dtype) and returns: int