ValueError (dimension off by 1) for HMC on Dirichlet

[jackkamm]

I am trying to do HMC on a model with Dirichlet random variables. However, I always get an error message like,

ValueError: Dimensions must be equal, but are N and N+1 for ... with input shapes: [N], [N+1]

where N is the number of categories of the Dirichlet variable.

Here is a minimal example:

import numpy as np
import tensorflow as tf
import edward as ed
from edward import models

n = 10
sigma = .1

dirichlet = models.Dirichlet(concentration=tf.ones(n))
noise = sigma * models.Normal(loc=tf.zeros(n), scale=tf.ones(n))
noisy_dirichlet = dirichlet + noise

data = np.random.dirichlet(np.ones(n)) + sigma * np.random.normal(size=n)

n_mcmc_steps = 1000
mcmc_vars = {dirichlet: models.Empirical(tf.Variable(1.0/n*np.ones((n_mcmc_steps, n)), dtype=np.float32))}

inference = ed.HMC(mcmc_vars, data={noisy_dirichlet: data})
inference.initialize()
inference.run()

And the error I get in this case is:

ValueError: Dimensions must be equal, but are 10 and 11 for 'inference/sample/Dirichlet_1/log_prob/mul' (op: 'Mul') with input shapes: [10], [11].

[colehawkins]

I've run into the same problem. When I run @jackkamm's code I get the errors below.

It seems like Issues #785 and #808 are related, as is this discourse thread,. I've run the Beta-Bernoulli model from the tests given in #808 with no problems, and I'm running version 1.3.5 of Edward and version 1.7.0 of Tensorflow. @davmre do you have any advice?

File "/Users/colehawkins/Dropbox/python/Test Scripts/untitled5.py", line 19, in <module>
    inference.initialize()

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/edward/inferences/hmc.py", line 64, in initialize
    return super(HMC, self).initialize(*args, **kwargs)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/edward/inferences/monte_carlo.py", line 101, in initialize
    self.train = self.build_update()

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/edward/inferences/hmc.py", line 104, in build_update
    self.n_steps)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/edward/inferences/hmc.py", line 207, in leapfrog
    grad_log_joint = tf.gradients(log_joint(z_new), list(six.itervalues(z_new)))

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/edward/inferences/hmc.py", line 167, in _log_joint_unconstrained
    return self._log_joint(z_sample_transformed) + log_det_jacobian

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/edward/inferences/hmc.py", line 193, in _log_joint
    log_joint += tf.reduce_sum(z_copy.log_prob(dict_swap[z]))

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/distributions/distribution.py", line 716, in log_prob
    return self._call_log_prob(value, name)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/distributions/distribution.py", line 698, in _call_log_prob
    return self._log_prob(value, **kwargs)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/distributions/util.py", line 1330, in _fn
    return fn(*args, **kwargs)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/distributions/dirichlet.py", line 205, in _log_prob
    return self._log_unnormalized_prob(x) - self._log_normalization()

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/distributions/dirichlet.py", line 213, in _log_unnormalized_prob
    return math_ops.reduce_sum((self.concentration - 1.) * math_ops.log(x), -1)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py", line 971, in binary_op_wrapper
    return func(x, y, name=name)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py", line 1198, in _mul_dispatch
    return gen_math_ops.mul(x, y, name=name)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/gen_math_ops.py", line 4689, in mul
    "Mul", x=x, y=y, name=name)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
    op_def=op_def)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3292, in create_op
    compute_device=compute_device)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3332, in _create_op_helper
    set_shapes_for_outputs(op)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2496, in set_shapes_for_outputs
    return _set_shapes_for_outputs(op)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2469, in _set_shapes_for_outputs
    shapes = shape_func(op)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2399, in call_with_requiring
    return call_cpp_shape_fn(op, require_shape_fn=True)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/common_shapes.py", line 627, in call_cpp_shape_fn
    require_shape_fn)

  File "/Users/colehawkins/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/common_shapes.py", line 691, in _call_cpp_shape_fn_impl
    raise ValueError(err.message)

ValueError: Dimensions must be equal, but are 10 and 11 for 'inference/sample/Dirichlet/log_prob/mul' (op: 'Mul') with input shapes: [10], [11].

[jkafrouni]

Any update on this?