InvalidArgumentError: Shapes of all inputs must match: values[0].shape = [3,3,3,128,64] != values[1].shape = [3,3,3] [Op:Pack] name: loc

I am getting an InvalidArgumentError: Shapes of all inputs must match: values[0].shape = [3,3,3,128,64] != values[1].shape = [3,3,3] [Op:Pack] name: loc

I have no idea where the error is. Can someone please tell me where is the mistake? Here is the code:

def multivariate_normal_gamma_precision_fn():

    class InverseGammaLogProb(tf.keras.regularizers.Regularizer):

        def __init__(self, prior_a=1., prior_b=10.):
            self.dist = InverseGamma(concentration=prior_a, scale=prior_b)

        def __call__(self, x):
            regularization = -tf.reduce_sum(self.dist.log_prob(x))
            return regularization

    def _fn(dtype, shape, name, trainable, add_variable_fn):
        log_alphas = add_variable_fn(
            name=name + '_log_alphas',
            shape=[1, 1, 1], #[1, 1]
            initializer=tf.keras.initializers.zeros(),
            regularizer=InverseGammaLogProb(),
            constraint=None,
            dtype=dtype,
            trainable=trainable)

        log_alphas_tiled = tf.tile(log_alphas, [shape[0], shape[1], shape[2]]) 
        dist = MultivariateNormalLogDiag(tf.zeros(shape, dtype=dtype), log_alphas_tiled) 
        batch_ndims = tf.size(dist.batch_shape_tensor())
        final_dist = tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)
        return final_dist
    return _fn

class InverseGamma(tfp.distributions.InverseGamma):
    """InverseGamma distribution where the log_prob can be evaluated with a log_x value, avoids doing log(exp(log_x))
       to get the log(x) value needed for the log_prob """

    def _log_prob(self, log_x):
        return self._log_unnormalized_prob(log_x) - self._log_normalization()

    def _log_unnormalized_prob(self, log_x):
        return -(self.concentration + 1.) * log_x - self.scale / tf.exp(log_x)

Start by reproducing the supplied Python code and inspect the shapes passed to tf.tile and MultivariateNormalLogDiag inside _fn. Compare the resulting tensor shapes with the two shapes in the InvalidArgumentError, then document the specific input-shape mismatch and verify the corrected behavior by rerunning the example.