code: def fit_bnn_map(
      net: bnn.BNN,
      seed: jax.Array,
      x_train: jax.Array,
      y_train: jax.Array,
      num_particles: int = 8,
      **optimizer_kwargs,
  ) -> Tuple[PyTree, dict[str, jax.Array]]:
      """Use a MAP estimator to fit a BNN."""
      optimizer_kwargs["num_particles"] = num_particles
      model_seed, optimization_seed = jax.random.split(seed)
      model = _make_bayeux_model(
          net, model_seed, x_train, y_train, num_particles
      )
      res = model.optimize.optax_adam(
          seed=optimization_seed, **optimizer_kwargs
      )  # pytype: disable=attribute-error
      params = res.params
      loss = res.loss
      return params, {"loss": loss}
  
imports: import functools
  from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
  import bayeux
  import jax
  import jax.numpy
  from jaxtyping import PyTree
  import matplotlib.pyplot
  import numpy
  import pandas
  from tensorflow_probability.python.experimental.autobnn import bnn
  from tensorflow_probability.python.experimental.autobnn import util
  from tensorflow_probability.python.experimental.timeseries import metrics
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 2 for calling a global function with example usage: fit_bnn_map(net, seed, x_train, y_train, num_particles) and returns: params,  dict