code: def fit_bnn_vi(
      net: bnn.BNN,
      seed: jax.Array,
      x_train: jax.Array,
      y_train: jax.Array,
      batch_size: int = 16,
      num_draws: int = 128,
      **vi_kwargs,
  ) -> Tuple[PyTree, dict[str, jax.Array]]:
      """Use a MAP estimator to fit a BNN."""
      vi_kwargs["batch_size"] = batch_size
      vi_kwargs["num_samples"] = num_draws
      model_seed, vi_seed, draw_seed = jax.random.split(seed, num=3)
      model = _make_bayeux_model(
          net, model_seed, x_train, y_train, batch_size, for_vi=True
      )
      (
          surrogate_dist,
          loss,
      ) = model.vi.tfp_factored_surrogate_posterior(  # pytype: disable=attribute-error
          seed=vi_seed, **vi_kwargs
      )
      params = surrogate_dist.sample(
          seed=draw_seed, sample_shape=num_draws
      )
  
      params = jax.tree_map(
          lambda x: x.reshape((-1,) + x.shape[2:]), params
      )
      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 4 for calling a global function with example usage: fit_bnn_vi(net, seed, x_train, y_train, batch_size, num_draws) and returns: params,  dict