def METHOD_NAME(self, x_start, tol):
    """Tests that adam optimizer takes one and two steps correctly
    for univariate functions."""
    stepsize, gamma, delta = 0.1, 0.5, 0.8
    adam_opt = AdamOptimizer(stepsize, beta1=gamma, beta2=delta)
    univariate_funcs = [np.sin, lambda x: np.exp(x / 10.0), lambda x: x**2]
    grad_uni_fns = [
        lambda x: (np.cos(x),),
        lambda x: (np.exp(x / 10.0) / 10.0,),
        lambda x: (2 * x,),
    ]
    for gradf, f in zip(grad_uni_fns, univariate_funcs):
        adam_opt.reset()
        x_onestep = adam_opt.step(f, x_start)
        adapted_stepsize = stepsize * np.sqrt(1 - delta) / (1 - gamma)
        firstmoment = (1 - gamma) * gradf(x_start)[0]
        secondmoment = (1 - delta) * gradf(x_start)[0] * gradf(x_start)[0]
        x_onestep_target = x_start - adapted_stepsize * firstmoment / (
            np.sqrt(secondmoment) + 1e-8
        )
        assert np.allclose(x_onestep, x_onestep_target, atol=tol)
        x_twosteps = adam_opt.step(f, x_onestep)
        adapted_stepsize = stepsize * np.sqrt(1 - delta**2) / (1 - gamma**2)
        firstmoment = gamma * firstmoment + (1 - gamma) * gradf(x_onestep)[0]
        secondmoment = (
            delta * secondmoment + (1 - delta) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
        )
        x_twosteps_target = x_onestep - adapted_stepsize * firstmoment / (
            np.sqrt(secondmoment) + 1e-8
        )
        assert np.allclose(x_twosteps, x_twosteps_target, atol=tol)