Implement log_abs_det_jacobian for the SoftmaxTransform?

[rachtsingh]

(Maybe a bad idea)

See #113

[fritzo]

This is tricky because SoftmaxTransform would be used in two important ways:

1. (the existing way) is to project an n-dimensional unnormalized vector of logits down to an n-1-dimensional subspace (the simplex). This is very useful for parameterizing simplex-constrained distributions in optimization settings, e.g. parameterizing a Categorical distribution in variational inference or MAP estimation. In this use, SoftmaxTransform is not bijective (.bijective == False) and only guarantees the pseudoinverse equations t(t.inv(t(x))) == t(x) and t.inv(t(t.inv(y))) == t.inv(y).
2. (the new proposed way) is to biject normalized logit vectors to normalized probability vectors in the simplex. In this usage, SoftmaxTransform is bijective (.bijective == True) and can define a .log_abs_det_jacobian(). This is useful in TransformedDistributions.

Question: Can we support both use cases?

[fritzo]

I'm not sure it is even possible to implement RelaxedCategorical from TransformedDistribution(Gumbel(...), BotlzmannTransform()) because the transform would need to project out / marginalize over one coordinate. This is different behavior from any of our existing TransformedDistributions and from Tensorflow's transformed distributions. It is a research question whether this is meaningful and what would be needed to support it. For example, .rsample() would work as expected, but .log_prob() would not work. We could possibly implement this by overriding .log_prob() behavior to not rely on SoftmaxTransform.log_abs_det_jacobian()

class RelaxedOneHotCategorical(TransformedDistribution):
    def __init__(self, ...):
        super(RelaxedOneHotCategorical, self).__init__(Gumbel(...), SoftmaxTransform())
    def log_prob(self, value):
        return ...a custom computation...