# Path: optax/_src/base.py
# NO_PARAMS_MSG = (
#     'You are using a transformation that requires the current value of '
#     'parameters, but you are not passing `params` when calling `update`.')
# class TransformInitFn(typing_extensions.Protocol):
# class TransformUpdateFn(typing_extensions.Protocol):
# class GradientTransformation(NamedTuple):
# class EmptyState(NamedTuple):
#   def __call__(self, params: Params) -> OptState:
#   def __call__(
#       self,
#       updates: Updates,
#       state: OptState,
#       params: Optional[Params] = None
#     ) -> Tuple[Updates, OptState]:
# def identity() -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def set_to_zero() -> GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def stateless(
#     f: Callable[[Updates, Optional[Params]], Updates],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def stateless_with_tree_map(
#     f: Callable[[chex.Array, Optional[chex.Array]], chex.Array],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):


from typing import Any, Callable, Sequence, Tuple
from optax._src import base
import chex
import jax
import jax.numpy as jnp
    # since we do a Taylor expansion of function around the mean.
    grads = jax.grad(function)(mean_dist)
    hessians = jax.hessian(function)(mean_dist)
    assert hessians.ndim == 2
    control_variate = function(mean_dist)
    control_variate += jnp.dot(centered_sample, grads)
    control_variate += jnp.dot(
        jnp.dot(centered_sample, hessians), centered_sample) / 2.
    return control_variate

  def expected_value_delta(
      params: base.Params, state: CvState) -> float:
    """"Expected value of second order expansion of `function` at dist mean."""
    del state
    mean_dist = params[0]
    var_dist = jnp.square(jnp.exp(params[1]))
    hessians = jax.hessian(function)(mean_dist)

    assert hessians.ndim == 2
    hess_diags = jnp.diag(hessians)
    assert hess_diags.ndim == 1

    # Trace (Hessian * Sigma) and we use that Sigma is diagonal.
    expected_second_order_term = jnp.sum(var_dist * hess_diags) / 2.

    expected_control_variate = function(mean_dist)
    expected_control_variate += expected_second_order_term
    return expected_control_variate

  def update_state(

# Path: optax/_src/base.py
# NO_PARAMS_MSG = (
#     'You are using a transformation that requires the current value of '
#     'parameters, but you are not passing `params` when calling `update`.')
# class TransformInitFn(typing_extensions.Protocol):
# class TransformUpdateFn(typing_extensions.Protocol):
# class GradientTransformation(NamedTuple):
# class EmptyState(NamedTuple):
#   def __call__(self, params: Params) -> OptState:
#   def __call__(
#       self,
#       updates: Updates,
#       state: OptState,
#       params: Optional[Params] = None
#     ) -> Tuple[Updates, OptState]:
# def identity() -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def set_to_zero() -> GradientTransformation:
#   def init_fn(params):
#   def update_fn(updates, state, params=None):
# def stateless(
#     f: Callable[[Updates, Optional[Params]], Updates],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):
# def stateless_with_tree_map(
#     f: Callable[[chex.Array, Optional[chex.Array]], chex.Array],
# ) -> GradientTransformation:
#   def init_fn(_):
#   def update_fn(updates, state, params=None):


# Path: optax/_src/control_variates.py
from typing import Any, Callable, Sequence, Tuple
from optax._src import base
import chex
import jax
import jax.numpy as jnp

    # since we do a Taylor expansion of function around the mean.
    grads = jax.grad(function)(mean_dist)
    hessians = jax.hessian(function)(mean_dist)
    assert hessians.ndim == 2
    control_variate = function(mean_dist)
    control_variate += jnp.dot(centered_sample, grads)
    control_variate += jnp.dot(
        jnp.dot(centered_sample, hessians), centered_sample) / 2.
    return control_variate

  def expected_value_delta(
      params: base.Params, state: CvState) -> float:
    """"Expected value of second order expansion of `function` at dist mean."""
    del state
    mean_dist = params[0]
    var_dist = jnp.square(jnp.exp(params[1]))
    hessians = jax.hessian(function)(mean_dist)

    assert hessians.ndim == 2
    hess_diags = jnp.diag(hessians)
    assert hess_diags.ndim == 1

    # Trace (Hessian * Sigma) and we use that Sigma is diagonal.
    expected_second_order_term = jnp.sum(var_dist * hess_diags) / 2.

    expected_control_variate = function(mean_dist)
    expected_control_variate += expected_second_order_term
    return expected_control_variate

  def update_state(

      params: base.Params,