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.

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