def METHOD_NAME(self, x_start, tol):
    """Tests that momentum optimizer takes one and two steps correctly
    for univariate functions."""
    stepsize, gamma = 0.1, 0.5
    mom_opt = MomentumOptimizer(stepsize, momentum=gamma)
    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):
        mom_opt.reset()
        x_onestep = mom_opt.step(f, x_start)
        x_onestep_target = x_start - gradf(x_start)[0] * stepsize
        assert np.allclose(x_onestep, x_onestep_target, atol=tol)
        x_twosteps = mom_opt.step(f, x_onestep)
        momentum_term = gamma * gradf(x_start)[0]
        x_twosteps_target = x_onestep - (gradf(x_onestep)[0] + momentum_term) * stepsize
        assert np.allclose(x_twosteps, x_twosteps_target, atol=tol)