code: def test_sparse():
      for dtype in _dtypes_to_test(
          use_gpu=test_utils.is_gpu_available()
      ):
          # Initialize tf for numpy implementation.
          m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
          var0_np = np.array(
              [1.0, 1.0, 2.0], dtype=dtype.as_numpy_dtype
          )
          grads0_np = np.array(
              [0.1, 0.0, 0.1], dtype=dtype.as_numpy_dtype
          )
          var1_np = np.array(
              [3.0, 3.0, 4.0], dtype=dtype.as_numpy_dtype
          )
          grads1_np = np.array(
              [0.01, 0.0, 0.01], dtype=dtype.as_numpy_dtype
          )
  
          var0 = tf.Variable(var0_np)
          var1 = tf.Variable(var1_np)
          grads0_np_indices = np.array([0, 2], dtype=np.int32)
          grads0 = tf.IndexedSlices(
              tf.constant(grads0_np[grads0_np_indices]),
              tf.constant(grads0_np_indices),
              tf.constant([3]),
          )
          grads1_np_indices = np.array([0, 2], dtype=np.int32)
          grads1 = tf.IndexedSlices(
              tf.constant(grads1_np[grads1_np_indices]),
              tf.constant(grads1_np_indices),
              tf.constant([3]),
          )
          opt = lamb.LAMB()
  
          # Fetch params to validate initial values
          np.testing.assert_allclose(
              np.asanyarray([1.0, 1.0, 2.0]), var0.numpy()
          )
          np.testing.assert_allclose(
              np.asanyarray([3.0, 3.0, 4.0]), var1.numpy()
          )
  
          # Run 3 steps of LAMB
          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 = lamb_update_numpy(
                  var0_np, grads0_np, t, m0, v0
              )
              var1_np, m1, v1 = lamb_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()
              )
  
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 4 for calling a global function with example usage: test_sparse() without return types