def METHOD_NAME(self, tol):
    """Tests that momentum optimizer takes one and two steps correctly
    for multivariate functions."""
    stepsize, gamma = 0.1, 0.5
    mom_opt = MomentumOptimizer(stepsize, momentum=gamma)
    multivariate_funcs = [
        lambda x: np.sin(x[0]) + np.cos(x[1]),
        lambda x: np.exp(x[0] / 3) * np.tanh(x[1]),
        lambda x: np.sum([x_**2 for x_ in x]),
    ]
    grad_multi_funcs = [
        lambda x: (np.array([np.cos(x[0]), -np.sin(x[1])]),),
        lambda x: (
            np.array(
                [
                    np.exp(x[0] / 3) / 3 * np.tanh(x[1]),
                    np.exp(x[0] / 3) * (1 - np.tanh(x[1]) ** 2),
                ]
            ),
        ),
        lambda x: (np.array([2 * x_ for x_ in x]),),
    ]
    x_vals = np.linspace(-10, 10, 16, endpoint=False)
    for gradf, f in zip(grad_multi_funcs, multivariate_funcs):
        for jdx in range(len(x_vals[:-1])):
            mom_opt.reset()
            x_vec = x_vals[jdx : jdx + 2]
            x_onestep = mom_opt.step(f, x_vec)
            x_onestep_target = x_vec - gradf(x_vec)[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_vec)[0]
            x_twosteps_target = x_onestep - (gradf(x_onestep)[0] + momentum_term) * stepsize
            assert np.allclose(x_twosteps, x_twosteps_target, atol=tol)