Issues with csiszar divergences

I was doing some work of mine regarding the f-divergences, and was happy to find their implementations in TFP, however I feel like there're several issues at the moment.

1. The tfp.vi.jensen_shannon divergence by default (with self_normalized=False) can be negative, since for f(u) = u log(u) - (1 + u) log(1 + u) f(1) = -2 ln(2). I don't see why one needs self_normalized at all, since f(u) = u log(u) - (1 + u) log(1 + u) + (u + 1) log(2) already satisfies both f(u=1)=0 and f'(u=1)=0. Same for the tfp.vi.arithmetic_geometric. Also, current implementation of the JSD is numerically unstable, but doesn't have to be. Here's a graph comparing current f(logu) implementation against mine (with a control variate built in, see the next point):

   def _f_jsd_stable(x):
       lsx = tf.math.log_sigmoid(x)
       return tf.where(lsx != 0., tf.exp(x) * lsx, -tf.ones_like(x)) + lsx - x + 2. * tf.log(2.)
   

2. tfp.vi.total_variation, tfp.vi.triangular and tfp.vi.jensen_shannon divergences would benefit from a control variate of the form a (u-1). In particular, for the Total Variation one can use 0.5 (u-1) as a control variate to obtain f(u) = max(0, 1-u), which is amazing since 1) it gives the same expectation, and thus defines the same divergence (for normalised measures); 2) this way f(u) ∈ [0, 1] greatly reducing the variance; 3) all operations can be made numerically stable. Similar trick can be pulled with the Triangle divergence, leading to the f(u) = 4 / (1+u) - 2

   Other divergences do not seem to be amendable by such control variate, and I couldn't come up with any practical alternative. The only possibility is c(u) = (a u + b) log u which can be expressed E_q(x) c(p(x)/q(x)) = a KL[p(x) || q(x)] - b KL[q(x) || p(x)], but this is only helpful for distributions with tractable KL, and is hard to implement in the current API.

I would be happy to contribute fixes

Start with the tfp.vi.jensen_shannon, tfp.vi.arithmetic_geometric, tfp.vi.total_variation, and tfp.vi.triangular divergence implementations. Compare their formulas with the proposed normalized and control-variate forms, then check numerical behavior using the _f_jsd_stable example. Done means the affected divergences are nonnegative where expected and numerically stable without changing their intended expectations.