Open martin-wiebusch-thg opened 3 weeks ago
The error occurs here:
File ~/.local/opt/miniconda/envs/adhoc/lib/python3.11/site-packages/tensorflow_probability/python/mcmc/internal/leapfrog_integrator.py:291, in SimpleLeapfrogIntegrator.call(self, momentum_parts, state_parts, target, target_grad_parts, kinetic_energy_fn, name)
an these seem to be the relevant parts of the message:
ValueError: in user code: ... File "/root/.local/opt/miniconda/envs/adhoc/lib/python3.11/site-packages/tensorflow_probability/python/mcmc/internal/leapfrog_integrator.py", line 291, in call ] = tf.while_loop(
ValueError: The two structures don't have the same nested structure.
... More specifically: Substructure "type=IndexedSlices str=IndexedSlices(indices=Tensor("mcmc_sample_chain/trace_scan/while/smart_for_loop/while/simple_step_size_adaptation_init/_one_step/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/gradients/mcmc_sample_chain/trace_scan/while/smart_for_loop/while/simple_step_size_adaptation_init/_one_step/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/GatherV2_grad/Reshape_1:0", shape=(3,), dtype=int32), values=Tensor("mcmc_sample_chain/trace_scan/while/smart_for_loop/while/simple_step_size_adaptation_init/_one_step/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/gradients/mcmc_sample_chain/trace_scan/while/smart_for_loop/while/simple_step_size_adaptation_init/_one_step/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/GatherV2_grad/Reshape:0", shape=(3,), dtype=float32), dense_shape=Tensor("mcmc_sample_chain/trace_scan/while/smart_for_loop/while/simple_step_size_adaptation_init/_one_step/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/gradients/mcmc_sample_chain/trace_scan/while/smart_for_loop/while/simple_step_size_adaptation_init/_one_step/mh_one_step/hmc_kernel_one_step/leapfrog_integrate/while/leapfrog_integrate_one_step/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/GatherV2_grad/Cast:0", shape=(1,), dtype=int32))" is a sequence, while substructure "type=SymbolicTensor str=Tensor("mcmc_sample_chain/trace_scan/while/smart_for_loop/while/simple_step_size_adaptation_init/_one_step/mh_one_step/hmc_kernel_one_step/maybe_call_fn_and_grads/value_and_gradients/fn_grad:0", shape=(None,), dtype=float32)" is not Entire first structure: [., [.], [.], ., [.]] Entire second structure: [., [.], [.], ., [.]]
I am trying to run a Hamiltonian MCMC on a target distribution whose implementation involves a call to tf.gather. The following code:
raises
ValueError: The two structures don't have the same nested structure.
followed by a very long and (to me) cryptic message. Replacing the return statement in the logprob function with the commented line gets rid of the error. The error seems to appear whenever the result of logprob contains atf.gather
subexpression.The error also disappears when I remove the
@tf.function
decorator from the definition ofrun_chain
. However, this comes at a huge performance cost.How can I efficiently sample from a distribution whose log-probability involves a
tf.gather
expression?