docs: def bayes_by_backprop_getter(
      posterior_builder=diagonal_gaussian_posterior_builder,
      prior_builder=fixed_gaussian_prior_builder,
      kl_builder=stochastic_kl_builder,
      sampling_mode_tensor=None,
      fresh_noise_per_connection=True,
      keep_control_dependencies=False):
    """Creates a custom getter which does Bayes by Backprop.
  
    Please see `tf.get_variable` for general documentation on custom getters.
  
    All arguments are optional. If nothing is configued, then a diagonal gaussian
    posterior will be used, and a fixed N(0, 0.01) prior will be used. Please
    see the default `posterior_builder` and `prior_builder` for a more detailed
    understanding of the default settings.
  
    Args:
      posterior_builder: A builder function which constructs an instance of
        `tfp.distributions.Distribution` which shall serve as the posterior over
        the `tf.Variable` of interest. The builder receives the `getter` and the
        arguments forwarded from `tf.get_variable`. Suppose one wrote
  
        ```
        tf.get_variable(
            'weights', shape=(3,), initializer=tf.zeros_initializer,
            dtype=tf.float32)
        ```
  
        then the `posterior_builder` argument would receive the `name`, `shape`,
        `initializer`, and `dtype` arguments passed above. The builder must return
        a `tfp.distributions.Distribution` object.
  
        Please see the `tf.get_variable` for documentation on `custom_getter` and
        `getter`, and see `bbb.diagonal_gaussian_posterior_builder`
        (the default) for an example of using this builder API.
      prior_builder: A builder function which constructs an instance of
        `tfp.distributions.Distribution` which shall serve as the prior over the
        `tf.Variable` of interest. Identical API to `posterior_builder`. See
        `bbb.fixed_gaussian_prior_builder` (the default) for an example.
      kl_builder: A builder function which receives the posterior distribution,
        prior distribution, and a sample from the posterior. It returns a
        scalar-shaped `tf.Tensor` representing the total KL cost for the
        `tf.Variable` in question. See `bbb.stochastic_kl_builder` (default) and
        `bbb.analytic_kl_builder` for examples.
      sampling_mode_tensor: A `tf.Tensor` which determines how an estimate from
        the posterior is produced. It must be scalar-shaped and have a `dtype` of
        `tf.string`. Valid values for this tensor are `bbb.EstimatorModes.sample`
        (which is the default), `bbb.EstimatorModes.mean`, and
        `bbb.EstimatorModes.last_sample`. `bbb.EstimatorModes.sample` is
        appropriate for training, and `bbb.EstimatorModes.mean` can be used
        at test time.
      fresh_noise_per_connection: A boolean. Indicates that each time a stochastic
        variable is retrieved with this custom getter, new sampling noise should
        be used. This is `True` by default. If this argument is set to `False`,
        then the same noise is used for each connection. Note that this does not
        apply to connections within a `tf.while_loop`; the same sampling noise
        is always used in different iterations of a `tf.while_loop` within one
        `session.run()` call. See the unit tests for details.
      keep_control_dependencies: A boolean. This argument should only be
        used by advanced users. Indicates that each time a stochastic variable is
        retrieved in the loop body of a `tf.while_loop` construct, new sampling
        noise should be used.
        The default behavior is `False`, so that RNNs use the same weights at each
        recurrent time step. This is done by removing the creation of the Variable
        from any existing control flow contexts. Notably, the Variables will be
        created outside the context of any tf.while_loop, making them fetchable.
        When this argument is `True`, any Variables used in the loop body of a
        `tf.while_loop` will be non-fetchable. If the KL cost needs to be
        evaluated, the Variable must *first* be used *outside* the loop body. This
        op using the Variable simply needs to be placed on the graph to get a
        stochastic estimate of the KL; it doesn't need to ever be used. Example:
  
        ```
        def loop_body(i):
          logits = sonnet_module(queue)
          i = i + 1
  
        with tf.variable_scope('bbb', custom_getter=bbb.bayes_by_backprop_getter(
            fresh_noise_per_connection=True,
            keep_control_dependencies=True)):
          unused_op = sonnet_module(queue)  # Adds KL estimate to bbb Collection
          final_i = tf.while_loop(lambda i: i < 5, loop_body, tf.constant(0.))
        ```
  
        Here when we add `unused_op` to the graph, we also add a number of tensors
        associated with the particular stochastic variable, including its
        contribution to the KL cost, to a graph-level registry. These are
        organized in a per-stochastic-variable data structure and be accessed with
        `bbb.get_variable_metadata()`. Without this line, these Tensors would
        instead be added the first time the Variable is used in the while_loop,
        which would make them non-fetchable.
  
        In all cases, the KL cost is only added once per Variable, which is the
        correct behavior, since if a variable is used multiple times in a model,
        the KL cost should remain unaffected.
  
    Returns:
      A `custom_getter` function which implements Bayes by Backprop.
    """
  
    if sampling_mode_tensor is None:
      sampling_mode_tensor = tf.constant(EstimatorModes.sample)
  
    def custom_getter(getter, name, *args, **kwargs):
      """The custom getter that will be returned."""
      if not kwargs.get("trainable", True):
        return getter(name, *args, **kwargs)
      if kwargs["dtype"] not in _OK_DTYPES_FOR_BBB:
        raise ValueError("Disallowed data type {}.".format(kwargs["dtype"]))
  
      var_scope = tf.get_variable_scope()
      if var_scope.reuse and not fresh_noise_per_connection:
        # Re-use the sampling noise by returning the very same posterior sample
        # if configured to do so.
        the_match = [
            x for x in get_variable_metadata() if x.raw_variable_name == name]
        if not the_match:
          raise ValueError(
              "Internal error. No metadata for variable {}".format(name))
        if len(the_match) > 1:
          raise ValueError(
              "Multiple matches for variable {}. Matches: {}".format(
                  name, [x.raw_variable_name for x in the_match]))
  
        return the_match[0].posterior_estimate
  
      raw_variable_shape = kwargs["shape"]
  
      def construct_subgraph():
        """Constructs subgraph used to reparameterize the variable in question."""
        posterior = posterior_builder(
            getter,
            name=name,
            *args, **kwargs)
        prior = prior_builder(
            getter,
            name=name,
            *args, **kwargs)
  
        # If the user does not return an extra dictionary of prior variables,
        # then fill in an empty dictionary.
        if isinstance(posterior, collections.Sequence):
          posterior_dist, posterior_vars = posterior
        else:
          posterior_dist, posterior_vars = posterior, {}
  
        if isinstance(prior, collections.Sequence):
          prior_dist, prior_vars = prior
        else:
          prior_dist, prior_vars = prior, {}
  
        if posterior_dist.reparameterization_type != _OK_PZATION_TYPE:
          raise ValueError(
              "Distribution {} incompatible with Bayes by Backprop.".format(
                  posterior_dist.__class__.__name__))
  
        posterior_estimator = _produce_posterior_estimate(posterior_dist,
                                                          sampling_mode_tensor,
                                                          name)
        kl_cost = kl_builder(posterior_dist, prior_dist, posterior_estimator)
        variable_metadata = _VariableMetadata(
            raw_variable_name=name,
            raw_variable_shape=raw_variable_shape,
            scope_name=var_scope.name,
            posterior=posterior_dist,
            posterior_estimate=posterior_estimator,
            prior=prior_dist,
            kl_cost=kl_cost,
            prior_vars=prior_vars,
            posterior_vars=posterior_vars)
        return posterior_estimator, variable_metadata
  
      # Entering the `tf.control_dependencies(None)` context is crucial to
      # provide compatibility with `tf.while_loop` and thus RNNs. The main thing
      # it does is making the `kl_cost` fetchable by causing these ops to be
      # created outside the context of any tf.while_loop. Note also that it causes
      # a RNN core's weights to be sampled just once when unrolled over a
      # sequence, rather than at every timestep.
      control_deps = [] if keep_control_dependencies else None
      with tf.control_dependencies(control_deps):
        posterior_estimator, var_metadata = construct_subgraph()
  
      # Only add these ops to a collection once per unique variable.
      # This is to ensure that KL costs are not tallied up more than once.
      var_with_name = _all_var_metadata_registry[tf.get_default_graph()].get(name)
      if var_with_name is None:
        _all_var_metadata_registry[tf.get_default_graph()][name] = var_metadata
  
      return posterior_estimator
    return custom_getter