mcmc model run error - tensorflow probability

Question

I am creating a utility model to evaluate a set of items. I have created the model in tesorflow probabilty, I am able to sample the pooled_model() successfully and even run the utilities function successfully manually with the samples. The shape of the ouput of the manual run matches the shape of all_choices. But if i run the model using mcmc I am getting an error:

%%time

@tf.function
def utilities(x, betas, error):
    """
    `x * betas + errors` with broadcasting.
    """
    x = tf.cast(x, dtype=tf.float32)
    betas = tf.squeeze(betas)
    util = (x * betas) + error 
    return util 

k = 10
sigma_beta = 1.
sigma_error = 1.


def pooled_model(X_train):
    return tfd.JointDistributionSequential([
        tfd.HalfCauchy(loc=0., scale=sigma_beta, name="sigma_beta"), 
        tfd.HalfCauchy(loc=0., scale=sigma_error, name="sigma_error"),
        tfd.Normal(loc=tf.zeros(k), scale=sigma_beta, name="beta"), 
        tfd.Gumbel(loc=0., scale=sigma_error, name="error"),
        lambda error, beta: tfd.Deterministic(
                tf.concat([tf.math.argmax(
                    tfd.Multinomial(
                        total_count=1, 
                        logits=utilities(X_train, beta[..., tf.newaxis], error[..., tf.newaxis]),
                    ).sample(), axis=1
                ), tf.math.argmin(
                    tfd.Multinomial(
                        total_count=1, 
                        logits=utilities(X_train, beta[..., tf.newaxis], error[..., tf.newaxis]),
                    ).sample(), axis=1
                )], 0)
            ),
    ])


def target_log_prob(sigma_beta, sigma_error, beta, error):
    return pooled_model(X_train).log_prob(sigma_beta, sigma_error, beta, error, all_choices)

# Use NUTS for inference
hmc = tfp.mcmc.NoUTurnSampler(
    target_log_prob_fn=target_log_prob,
    step_size=.01)

ValueError: Dimensions must be equal, but are 10 and 120000 for '{{node mcmc_sample_chain/ /_bootstrap_results/transformed_kernel_bootstrap_results/NoUTurnSampler/.bootstrap_results/process_args/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/JointDistributionSequential/log_prob/add_4}} = AddV2[T=DT_FLOAT](mcmc_sample_chain/dual_averaging_step_size_adaptation___init__/_bootstrap_results/transformed_kernel_bootstrap_results/NoUTurnSampler/.bootstrap_results/process_args/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/JointDistributionSequential/log_prob/add_3, mcmc_sample_chain/dual_averaging_step_size_adaptation___init__/_bootstrap_results/transformed_kernel_bootstrap_results/NoUTurnSampler/.bootstrap_results/process_args/maybe_call_fn_and_grads/value_and_gradients/value_and_gradient/JointDistributionSequential/log_prob/Deterministic/log_prob/Log)' with input shapes: [1,10], [120000].

X_train is of the shape -> (60000, 10) all_choices is of the shape (120000,)

I am expecting to be able run the model successfully