'How to construct an equivalent multivariate normal distribution in tensorflow-probability, using TransformedDistribution?

How to construct an equivalent multivariate normal distribution in tensorflow-probability, using TransformedDistribution and tfb.ScaleMatvecLinearOperator?

I'm reading about a tutorial on a bijector in tensorflow_probability: tfp.bijectors.ScaleMatvecLinearOperator. An example was provided.

n = 10000
loc = 0
scale = 0.5
normal = tfd.Normal(loc=loc, scale=scale)

The above codes creates a univariate normal distribution.

tril = tf.random.normal((2, 4, 4)) 
scale_low_tri = tf.linalg.LinearOperatorLowerTriangular(tril)
scale_low_tri.to_dense()

The above codes created a tensor consisting of 2 lower triangular matrix:

<tf.Tensor: shape=(2, 4, 4), dtype=float32, numpy=
array([[[-0.56953585,  0.        ,  0.        ,  0.        ],
        [ 1.1368589 ,  0.32028311,  0.        ,  0.        ],
        [-0.8328388 , -1.9963025 , -0.6005632 ,  0.        ],
        [ 0.596155  , -0.214932  ,  1.0988408 , -0.41731614]],

       [[ 2.0778096 ,  0.        ,  0.        ,  0.        ],
        [-1.1863967 ,  2.4897904 ,  0.        ,  0.        ],
        [ 0.38001925,  1.4962028 ,  1.7609248 ,  0.        ],
        [ 2.9253726 ,  0.7047957 ,  0.050508  ,  0.58643174]]],
      dtype=float32)>

Then a matrix-vector multiplication bijector is created:

scale_lin_op = tfb.ScaleMatvecLinearOperator(scale_low_tri)

After that, a TransformedDistribution is constructed as follows:

mvn = tfd.TransformedDistribution(normal, scale_lin_op, batch_shape=[2],  event_shape=[4]) # 

This should have worked in the old versions of tensorflow_probability. However the constructor of TransformedDistribution is changed now and does not accept the last two parameters batch_shape and event_shape. Therefore I tried to use the following way to do the same:

mvn2 = tfd.TransformedDistribution(
    distribution=tfd.Sample(
        normal,
        sample_shape=[4] # base_dist.event_shape == [4]
    ),
    bijector=scale_lin_op, ) # batch_shape=[2],  event_shape=[4]
mvn2

And the result seems to have the correct batch_shape and event_shape

<tfp.distributions.TransformedDistribution 'scale_matvec_linear_operatorSampleNormal' batch_shape=[2] event_shape=[4] dtype=float32>

Then, another distribution for comparison is created:

mvn3 = tfd.MultivariateNormalLinearOperator(loc=loc, scale=scale_low_tri)
mvn3

According to the tutorial, the TransformedDistribution mvn2 should be equivalent to the MultivariateNormalLinearOperator mvn3.

# Check
xn = normal.sample((n, 2, 4)) # sample_shape = (n, 2, 4)
tf.norm(mvn2.log_prob(xn) - mvn3.log_prob(xn)) / tf.norm(mvn2.log_prob(xn))

<tf.Tensor: shape=(), dtype=float32, numpy=0.7498207>

But in my result they are not equivalent. (If they are, the above tensor should be 0)

What have I done wrong?



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source