def METHOD_NAME(self, tol):
    """Tests that adam optimizer takes one and two steps correctly
    for multivariate functions."""
    stepsize, gamma, delta = 0.1, 0.5, 0.8
    adam_opt = AdamOptimizer(stepsize, beta1=gamma, beta2=delta)
    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])):
            adam_opt.reset()
            x_vec = x_vals[jdx : jdx + 2]
            x_onestep = adam_opt.step(f, x_vec)
            adapted_stepsize = stepsize * np.sqrt(1 - delta) / (1 - gamma)
            firstmoment = (1 - gamma) * gradf(x_vec)[0]
            secondmoment = (1 - delta) * gradf(x_vec)[0] * gradf(x_vec)[0]
            x_onestep_target = x_vec - 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)