Need of Exp bijector in the learning the Normal? I am trying to understand TransformedVariable class in tensorflow probability. The website provides the following example:
trainable_normal = tfd.Normal(
    loc=tf.Variable(0.),
    scale=tfp.util.TransformedVariable(1., bijector=tfb.Exp()))

with tf.GradientTape() as tape:
  negloglik = -trainable_normal.log_prob(0.5)
g = tape.gradient(negloglik, trainable_normal.trainable_variables)
opt = tf.optimizers.Adam(learning_rate=0.05)
loss = tf.function(lambda: -trainable_normal.log_prob(0.5))

for _ in range(int(1e3)):
  opt.minimize(loss, trainable_normal.trainable_variables)

print(trainable_normal.mean())
print(trainable_normal.stddev())

Where the variance is modeled using exp bijector for presumably restricting it to positive values.
If that is the case shouldn't any other positive semidefinite bijector should work the same?
However when I change the bijector to tensorflow_probability.bijectors.Square the answers are all wrong:
#  For Exp TransformedVariable
tf.Tensor(0.5, shape=(), dtype=float32)
tf.Tensor(0.00021710133, shape=(), dtype=float32

# For Square TransformedVariable
tf.Tensor(0.891498, shape=(), dtype=float32)
tf.Tensor(0.24519879, shape=(), dtype=float32)

Why is TransformedVariable needed here, why it should be Exp?
 A: The purpose of a bijector in TensorFlow Probability is ultimately used as an interface for transforming the distribution of a sample.
This is ultimately achieved by three methods as outlined at the TensorFlow website:

*

*Forward

*Inverse

*log_det_jacobian(x)
The guide cited gives more detail, but your choice of which one to use depends on your goals.
For instance, if you wished to transform the distribution of the variable in question into a different distribution by creating new samples - the forward method is used. The inverse is used for computing probabilities using the existing samples rather than creating new samples outright.
In the case of Exp vs Square, it is not particularly surprising that the results will be different as they are technically different distributions - even if they are both being used to transform the distribution into positive values.
Let's take a set of 100 normally distributed numbers. Let's call this array a.

Now, let us square the observations in array a (array b) and find the exponent of the observations in array a (array c).
Distribution for array b

Distribution for array c

We can see that while b and c assume the same shape, the distribution is more spread out for c.
In particular, if we calculate the standard deviation and mean for arrays a, b, and c, we find that c has the highest values across these parameters:
>>> np.std(a)
0.9568046362169994
>>> np.std(b)
1.3206259644496021
>>> np.std(c)
1.935180654408512
>>> np.mean(a)
0.10465525532719616
>>> np.mean(b)
0.9264278343539454
>>> np.mean(c)
1.7425434179206918

Therefore, I am unsure as to why one might expect to obtain the same results if the Squared distribution is being used instead of Exp. While both are transforming the values into positive ones, there are still inherent differences between the two distributions.
Hope this is of help.
