# 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
  Returns:
    A `num_samples x D` vector containing the estimates of the gradients
    obtained for each sample. The mean of this vector can be used to update
    the scale parameter. The entire vector can be used to assess estimator
    variance.
  """
  mean, log_std = dist.params
  std = jnp.exp(log_std)

  dist_samples = dist.sample((num_samples,), seed=rng)

  pos_rng, neg_rng = jax.random.split(rng)

  # The only difference between mean and std gradients is what we sample.
  pos_sample = jax.random.double_sided_maxwell(
      pos_rng, loc=0.0, scale=1.0, shape=dist_samples.shape)
  if coupling:
    unif_rvs = jax.random.uniform(neg_rng, dist_samples.shape)
    neg_sample = unif_rvs * pos_sample
  else:
    neg_sample = jax.random.normal(neg_rng, dist_samples.shape)

  # Both need to be positive in the case of the scale.
  # N x D
  positive_diag = mean + std * pos_sample
  # N x D
  negative_diag = mean + std * neg_sample

  # NOTE: you can sample base samples here if you use the same rng
  # Duplicate the D dimension - N x D x D.

# 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

  Returns:
    A `num_samples x D` vector containing the estimates of the gradients
    obtained for each sample. The mean of this vector can be used to update
    the scale parameter. The entire vector can be used to assess estimator
    variance.
  """
  mean, log_std = dist.params
  std = jnp.exp(log_std)

  dist_samples = dist.sample((num_samples,), seed=rng)

  pos_rng, neg_rng = jax.random.split(rng)

  # The only difference between mean and std gradients is what we sample.
  pos_sample = jax.random.double_sided_maxwell(
      pos_rng, loc=0.0, scale=1.0, shape=dist_samples.shape)
  if coupling:
    unif_rvs = jax.random.uniform(neg_rng, dist_samples.shape)
    neg_sample = unif_rvs * pos_sample
  else:
    neg_sample = jax.random.normal(neg_rng, dist_samples.shape)

  # Both need to be positive in the case of the scale.
  # N x D
  positive_diag = mean + std * pos_sample
  # N x D
  negative_diag = mean + std * neg_sample

  # NOTE: you can sample base samples here if you use the same rng
  # Duplicate the D dimension - N x D x D.

  base_dist_samples = utils.tile_second_to_last_dim(dist_samples)