MAP and Maximum-Likelihood as VI

I am using the following code to seamlessly switch between variational, MAP and ML approaches, something that I found useful for prototyping new variational inference ideas:

class ImproperUniform(tfd.Distribution):
    def __init__(self, shape=None, dtype=tf.float32, nevent_dims=1, validate_args=False, allow_nan_stats=True):
        super().__init__(
            dtype,
            reparameterization_type=tfd.FULLY_REPARAMETERIZED,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats)

        self.shape = shape
        self.batch_dims = len(self.shape) - nevent_dims

    def _prob(self, x):
        # correct up to a constant, ok if only the gradient matters:
        return tf.ones(tf.shape(x)[:self.batch_dims])

    def _log_prob(self, x):
        # correct up to a constant, ok if only the gradient matters:
        return tf.zeros(tf.shape(x)[:self.batch_dims])

    def _batch_shape_tensor(self):
        return self.shape[:self.batch_dims]

    def _batch_shape(self):
        return self.shape[:self.batch_dims]

    def _event_shape_tensor(self):
        return self.shape[self.batch_dims:]

    def _event_shape(self):
        return self.shape[self.batch_dims:]


@tfd.RegisterKL(tfd.Deterministic, tfd.Distribution)
def kl_deterministic_distribution(a: tfd.Deterministic, b: tfd.Distribution, name=None):
    with tf.name_scope(name, 'kl_deterministic_distribution', [a.loc, b.parameters]):
        # correct up to a constant, ok if only the gradient matters:
        return -b.log_prob(a.loc)


@tfd.RegisterKL(tfd.VectorDeterministic, tfd.Distribution)
def kl_vectorDeterministic_distribution(a: tfd.VectorDeterministic, b: tfd.Distribution, name=None):
    with tf.name_scope(name, 'kl_vectorDeterministic_distribution', [a.loc, b.parameters]):
        # correct up to a constant, ok if only the gradient matters:
        return -b.log_prob(a.loc)


@tfd.RegisterKL(tfd.Distribution, ImproperUniform)
def kl_distribution_improperUniform(a: tfd.Distribution, b: ImproperUniform, name=None):
    with tf.name_scope(name, 'kl_distribution_improperUniform', [a.parameters]):
        # correct up to a constant, ok if only the gradient matters:
        return -a.entropy()

What's unusual about this is that only the gradients of log_prop and KLs are correct, while the absolute value has a different meaning, e. g. the negative logprob under the prior in case of MAP instead of the KL. One could therefore also consider creating a new set of methods that make this semantic difference explicit, like "log_prob_plus_const" and "kl_divergence_plus_const", which uses the "log_prob" and registered KL methods by default but which can be augmented by the methods above.

In case this is something you consider worth integrating into TFP, I would potentially be interested in creating a PR in the future. Also happy to provide proofs for the correctness of the gradients.

Start by reviewing the proposed ImproperUniform distribution and the three registered KL functions in the issue body. Determine whether TFP should expose explicit methods for values that are only correct up to a constant, and define the API and validation needed before implementation. Done means an agreed integration design or a clearly scoped contribution.