code: def _make_bayeux_model(
      net: bnn.BNN,
      seed: jax.Array,
      x_train: jax.Array,
      y_train: jax.Array,
      num_particles: int = 8,
      for_vi: bool = False,
  ):
      """Use a MAP estimator to fit a BNN."""
      test_seed, init_seed = jax.random.split(seed)
      test_point = net.init(test_seed, x_train)
      transform, inverse_transform, ildj = util.make_transforms(net)
  
      def _init(seed):
          return net.init(seed, x_train)
  
      initial_state = jax.vmap(_init)(
          jax.random.split(init_seed, num_particles)
      )
  
      if for_vi:
  
          def log_density(params, *, seed=None):
              # The TFP VI machinery tries passing in a `seed` parameter, splats
              # dictionaries that are used for arguments, and adds a batch dimension
              # of size [1] to the start, so we undo all of that.
              del seed
              return net.log_prob(
                  {"params": jax.tree_map(lambda x: x[0, ...], params)},
                  data=x_train,
                  observations=y_train,
              )
  
      else:
          log_density = functools.partial(
              net.log_prob, data=x_train, observations=y_train
          )
      return bx.Model(
          log_density=log_density,
          test_point=test_point,
          initial_state=initial_state,
          transform_fn=transform,
          inverse_transform_fn=inverse_transform,
          inverse_log_det_jacobian=ildj,
      )
  
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 1 for calling a global function with example usage: _make_bayeux_model(net, seed, x_train, y_train, num_particles, for_vi) and returns: bx, net, net