code: def fit_bnn_mcmc(
      net: bnn.BNN,
      seed: jax.Array,
      x_train: jax.Array,
      y_train: jax.Array,
      num_chains: int = 128,
      num_draws: int = 8,
      **sampler_kwargs,
  ) -> Tuple[PyTree, dict[str, jax.Array]]:
      """Use a MAP estimator to fit a BNN."""
      sampler_kwargs["num_chains"] = num_chains
      sampler_kwargs["num_samples"] = num_draws
      sampler_kwargs["return_pytree"] = True
      model_seed, mcmc_seed = jax.random.split(seed)
      model = _make_bayeux_model(
          net, model_seed, x_train, y_train, num_chains
      )
      params = model.mcmc.numpyro_nuts(
          seed=mcmc_seed, **sampler_kwargs
      )  # pytype: disable=attribute-error
  
      # TODO(colcarroll): This function should instead reliably return `params``
      # with shape (num_chains * num_samples, ...), but looking at per-chain metrics
      # is the easiest way to determine where "stuck chains" occur, and it is
      # nice to return parameters with a single batch dimension.
      params = _filter_stuck_chains(params)
      params = jax.tree_map(
          lambda x: x.reshape((-1,) + x.shape[2:]), params
      )
      return params, {
          "noise_scale": params["params"].get("noise_scale", None)
      }
  
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 5 for calling a global function with example usage: fit_bnn_mcmc(net, seed, x_train, y_train, num_chains, num_draws) and returns: params,  dict