def METHOD_NAME(self):
    for dtype in [dtypes.half, dtypes.float32, dtypes.float64]:
        learning_rate = 0.0001
        beta_1 = 0.9
        beta_2 = 0.99
        with self.cached_session():
            m0_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
            m1_np = np.array([0.0, 0.0], dtype=dtype.as_numpy_dtype)
            var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
            var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
            grads0_np = np.array([0.9, 0.0], dtype=dtype.as_numpy_dtype)
            grads1_np = np.array([0.1, 0.0], dtype=dtype.as_numpy_dtype)
            var0 = tf.Variable(var0_np)
            var1 = tf.Variable(var1_np)
            grads0_np_indices = np.array([0], dtype=np.int32)
            grads0 = tf.IndexedSlices(
                tf.constant(grads0_np[grads0_np_indices]),
                tf.constant(grads0_np_indices),
                tf.constant([2]),
            )
            grads1_np_indices = np.array([0], dtype=np.int32)
            grads1 = tf.IndexedSlices(
                tf.constant(grads1_np[grads1_np_indices]),
                tf.constant(grads1_np_indices),
                tf.constant([2]),
            )
            optimizer = Lion(
                learning_rate=learning_rate,
                beta_1=beta_1,
                beta_2=beta_2,
            )
            # Run 3 steps of Lion
            for _ in range(3):
                optimizer.apply_gradients(
                    zip([grads0, grads1], [var0, var1])
                )
                var0_np, m0_np = lion_update_numpy(
                    var0_np,
                    grads0_np,
                    m0_np,
                    learning_rate=learning_rate,
                    beta_1=beta_1,
                    beta_2=beta_2,
                )
                var1_np, m1_np = lion_update_numpy(
                    var1_np,
                    grads1_np,
                    m1_np,
                    learning_rate=learning_rate,
                    beta_1=beta_1,
                    beta_2=beta_2,
                )
                # Validate updated params
                self.assertAllCloseAccordingToType(var0_np, var0)
                self.assertAllCloseAccordingToType(var1_np, var1)