# Path: optax/_src/second_order.py
# def ravel(p: Any) -> Array:
# def hvp(
#     loss: LossFun,
#     v: jnp.DeviceArray,
#     params: Any,
#     inputs: jnp.DeviceArray,
#     targets: jnp.DeviceArray,
# ) -> jnp.DeviceArray:
# def hessian_diag(
#     loss: LossFun,
#     params: Any,
#     inputs: jnp.DeviceArray,
#     targets: jnp.DeviceArray,
# ) -> jnp.DeviceArray:
# def fisher_diag(
#     negative_log_likelihood: LossFun,
#     params: Any,
#     inputs: jnp.ndarray,
#     targets: jnp.ndarray,
# ) -> jnp.DeviceArray:


import collections
import functools
import itertools
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from optax._src import second_order
    self.loss_fn = loss

    def jax_hessian_diag(loss_fun, params, inputs, targets):
      """This is the 'ground-truth' obtained via the JAX library."""
      hess = jax.hessian(loss_fun)(params, inputs, targets)

      # Extracts the diagonal components.
      hess_diag = collections.defaultdict(dict)
      for k0, k1 in itertools.product(params.keys(), ['w', 'b']):
        params_shape = params[k0][k1].shape
        n_params = np.prod(params_shape)
        hess_diag[k0][k1] = jnp.diag(hess[k0][k1][k0][k1].reshape(
            n_params, n_params)).reshape(params_shape)
      for k, v in hess_diag.items():
        hess_diag[k] = v
      return second_order.ravel(hess_diag)

    self.hessian = jax_hessian_diag(
        self.loss_fn, self.parameters, self.data, self.labels)

  @chex.all_variants()
  def test_hessian_diag(self):
    hessian_diag_fn = self.variant(
        functools.partial(second_order.hessian_diag, self.loss_fn))
    actual = hessian_diag_fn(self.parameters, self.data, self.labels)
    np.testing.assert_array_almost_equal(self.hessian, actual, 5)

  @chex.all_variants()
  def test_fisher_diag_shape(self):
    fisher_diag_fn = self.variant(

# Path: optax/_src/second_order.py
# def ravel(p: Any) -> Array:
# def hvp(
#     loss: LossFun,
#     v: jnp.DeviceArray,
#     params: Any,
#     inputs: jnp.DeviceArray,
#     targets: jnp.DeviceArray,
# ) -> jnp.DeviceArray:
# def hessian_diag(
#     loss: LossFun,
#     params: Any,
#     inputs: jnp.DeviceArray,
#     targets: jnp.DeviceArray,
# ) -> jnp.DeviceArray:
# def fisher_diag(
#     negative_log_likelihood: LossFun,
#     params: Any,
#     inputs: jnp.ndarray,
#     targets: jnp.ndarray,
# ) -> jnp.DeviceArray:


# Path: optax/_src/second_order_test.py
import collections
import functools
import itertools
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from optax._src import second_order

    self.loss_fn = loss

    def jax_hessian_diag(loss_fun, params, inputs, targets):
      """This is the 'ground-truth' obtained via the JAX library."""
      hess = jax.hessian(loss_fun)(params, inputs, targets)

      # Extracts the diagonal components.
      hess_diag = collections.defaultdict(dict)
      for k0, k1 in itertools.product(params.keys(), ['w', 'b']):
        params_shape = params[k0][k1].shape
        n_params = np.prod(params_shape)
        hess_diag[k0][k1] = jnp.diag(hess[k0][k1][k0][k1].reshape(
            n_params, n_params)).reshape(params_shape)
      for k, v in hess_diag.items():
        hess_diag[k] = v
      return second_order.ravel(hess_diag)

    self.hessian = jax_hessian_diag(
        self.loss_fn, self.parameters, self.data, self.labels)

  @chex.all_variants()
  def test_hessian_diag(self):
    hessian_diag_fn = self.variant(
        functools.partial(second_order.hessian_diag, self.loss_fn))
    actual = hessian_diag_fn(self.parameters, self.data, self.labels)
    np.testing.assert_array_almost_equal(self.hessian, actual, 5)

  @chex.all_variants()
  def test_fisher_diag_shape(self):
    fisher_diag_fn = self.variant(

        functools.partial(second_order.fisher_diag, self.loss_fn))