# Copyright Contributors to the Pyro project.
# SPDX-License-Identifier: Apache-2.0

from collections import namedtuple
import logging
import os

import numpy as np
from numpy.testing import assert_allclose
import pytest

from jax import device_put, disable_jit, grad, jit, random, tree_map
import jax.numpy as jnp

import numpyro.distributions as dist
from numpyro.infer.hmc_util import (
    AdaptWindow,
    _is_iterative_turning,
    _leaf_idx_to_ckpt_idxs,
    build_adaptation_schedule,
    build_tree,
    consensus,
    dual_averaging,
    find_reasonable_step_size,
    parametric_draws,
    velocity_verlet,
    warmup_adapter,
    welford_covariance,
)
from numpyro.util import control_flow_prims_disabled, fori_loop, optional

logger = logging.getLogger(__name__)


@pytest.mark.parametrize("jitted", [True, False])
def test_dual_averaging(jitted):
    def optimize(f):
        da_init, da_update = dual_averaging(gamma=0.5)