# 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/utils.py
# def tile_second_to_last_dim(a):
# def canonicalize_dtype(dtype):
# def cast_tree(tree, dtype):
# def set_diags(a, new_diags):
# def multi_normal(loc, log_scale):
#   def __init__(self, loc, log_scale):
#   def sample(self, shape, seed):
#   def log_prob(self, x):
#   def log_scale(self):
#   def params(self):
# def _scale_gradient(inputs: chex.Array, scale: float)  -> chex.ArrayTree:
# def _scale_gradient_fwd(inputs, scale):
# def _scale_gradient_bwd(scale, g):
# def scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
# class MultiNormalDiagFromLogScale():


import math
import chex
import jax
import jax.numpy as jnp
import numpy as np
from typing import Any, Callable, Sequence
from optax._src import base
from optax._src import utils
    function: Function f(x) for which to estimate grads_{params} E_dist f(x).
      The function takes in one argument (a sample from the distribution) and
      returns a floating point value.
    params: A tuple of jnp arrays.
      The parameters for which to construct the distribution.
    dist_builder: a constructor which builds a distribution given the input
      parameters specified by params. `dist_builder(params)` should return a
      valid distribution.
    rng: a PRNGKey key.
    num_samples: Int, the number of samples used to compute the grads.

  Returns:
    A tuple of size `params`, each element is `num_samples x param.shape`
      jacobian vector containing the estimates of the gradients obtained for
      each sample.
    The mean of this vector is the gradient wrt to parameters that can be used
      for learning. The entire jacobian vector can be used to assess estimator
      variance.
  """
  def surrogate(params):
    # We vmap the function application over samples - this ensures that the
    # function we use does not have to be vectorized itself.
    dist = dist_builder(*params)
    return jax.vmap(function)(dist.sample((num_samples,), seed=rng))

  return jax.jacfwd(surrogate)(params)


def measure_valued_jacobians(
    function: Callable[[chex.Array], float],

# 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/utils.py
# def tile_second_to_last_dim(a):
# def canonicalize_dtype(dtype):
# def cast_tree(tree, dtype):
# def set_diags(a, new_diags):
# def multi_normal(loc, log_scale):
#   def __init__(self, loc, log_scale):
#   def sample(self, shape, seed):
#   def log_prob(self, x):
#   def log_scale(self):
#   def params(self):
# def _scale_gradient(inputs: chex.Array, scale: float)  -> chex.ArrayTree:
# def _scale_gradient_fwd(inputs, scale):
# def _scale_gradient_bwd(scale, g):
# def scale_gradient(inputs: chex.ArrayTree, scale: float) -> chex.ArrayTree:
# class MultiNormalDiagFromLogScale():


# Path: optax/_src/stochastic_gradient_estimators.py
import math
import chex
import jax
import jax.numpy as jnp
import numpy as np
from typing import Any, Callable, Sequence
from optax._src import base
from optax._src import utils

    function: Function f(x) for which to estimate grads_{params} E_dist f(x).
      The function takes in one argument (a sample from the distribution) and
      returns a floating point value.
    params: A tuple of jnp arrays.
      The parameters for which to construct the distribution.
    dist_builder: a constructor which builds a distribution given the input
      parameters specified by params. `dist_builder(params)` should return a
      valid distribution.
    rng: a PRNGKey key.
    num_samples: Int, the number of samples used to compute the grads.

  Returns:
    A tuple of size `params`, each element is `num_samples x param.shape`
      jacobian vector containing the estimates of the gradients obtained for
      each sample.
    The mean of this vector is the gradient wrt to parameters that can be used
      for learning. The entire jacobian vector can be used to assess estimator
      variance.
  """
  def surrogate(params):
    # We vmap the function application over samples - this ensures that the
    # function we use does not have to be vectorized itself.
    dist = dist_builder(*params)
    return jax.vmap(function)(dist.sample((num_samples,), seed=rng))

  return jax.jacfwd(surrogate)(params)


def measure_valued_jacobians(
    function: Callable[[chex.Array], float],

    params: base.Params,