# Buggy code
def main(args):
    baseball_dataset = pd.read_csv(DATA_URL, "\t")
    train, _, player_names = train_test_split(baseball_dataset)
    at_bats, hits = train[:, 0], train[:, 1]
    nuts_kernel = NUTS(conditioned_model)
    logging.info("Original Dataset:")
    logging.info(baseball_dataset)

    # (1) Full Pooling Model
    posterior_fully_pooled = MCMC(
        nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
    ).run(fully_pooled, at_bats, hits)
    logging.info("\nModel: Fully Pooled")
    logging.info("===================")
    logging.info("\nphi:")
    logging.info(
        summary(posterior_fully_pooled, sites=["phi"], player_names=player_names)["phi"]
    )
    posterior_predictive = TracePredictive(
        fully_pooled, posterior_fully_pooled, num_samples=args.num_samples
    )
    sample_posterior_predictive(posterior_predictive, baseball_dataset)
    evaluate_log_predictive_density(
        fully_pooled, posterior_fully_pooled, baseball_dataset
    )

    # (2) No Pooling Model
    posterior_not_pooled = MCMC(
        nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
    ).run(not_pooled, at_bats, hits)
    logging.info("\nModel: Not Pooled")
    logging.info("=================")
    logging.info("\nphi:")
    logging.info(
        summary(posterior_not_pooled, sites=["phi"], player_names=player_names)["phi"]
    )
    posterior_predictive = TracePredictive(
        not_pooled, posterior_not_pooled, num_samples=args.num_samples
    )
    sample_posterior_predictive(posterior_predictive, baseball_dataset)
    evaluate_log_predictive_density(not_pooled, posterior_not_pooled, baseball_dataset)

    # (3) Partially Pooled Model
    # TODO: remove once htps://github.com/uber/pyro/issues/1458 is resolved
    if "CI" not in os.environ:
        posterior_partially_pooled = MCMC(
            nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
        ).run(partially_pooled, at_bats, hits)
        logging.info("\nModel: Partially Pooled")
        logging.info("=======================")
        logging.info("\nphi:")
        logging.info(
            summary(
                posterior_partially_pooled, sites=["phi"], player_names=player_names
            )["phi"]
        )
        posterior_predictive = TracePredictive(
            partially_pooled, posterior_partially_pooled, num_samples=args.num_samples
        )
        sample_posterior_predictive(posterior_predictive, baseball_dataset)
        evaluate_log_predictive_density(
            partially_pooled, posterior_partially_pooled, baseball_dataset
        )

    # (4) Partially Pooled with Logit Model
    posterior_partially_pooled_with_logit = MCMC(
        nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
    ).run(partially_pooled_with_logit, at_bats, hits)
    logging.info("\nModel: Partially Pooled with Logit")
    logging.info("==================================")
    logging.info("\nSigmoid(alpha):")
    logging.info(
        summary(
            posterior_partially_pooled_with_logit,
            sites=["alpha"],
            player_names=player_names,
            transforms={"alpha": lambda x: 1.0 / (1 + np.exp(-x))},
        )["alpha"]
    )
    posterior_predictive = TracePredictive(
        partially_pooled_with_logit,
        posterior_partially_pooled_with_logit,
        num_samples=args.num_samples,
    )
    sample_posterior_predictive(posterior_predictive, baseball_dataset)
    evaluate_log_predictive_density(
        partially_pooled_with_logit,
        posterior_partially_pooled_with_logit,
        baseball_dataset,
    )


# Traceback (KeyError)
Traceback (most recent call last):
File "/home/travis/build/uber/pyro/pyro/poutine/trace_struct.py", line 128, in log_prob_sum
log_p = site["log_prob_sum"]
KeyError: 'log_prob_sum'
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/home/travis/build/uber/pyro/examples/baseball.py", line 309, in <module>
main(args)
File "/home/travis/build/uber/pyro/examples/baseball.py", line 272, in main
.run(partially_pooled, at_bats, hits)
File "/home/travis/build/uber/pyro/pyro/infer/abstract_infer.py", line 84, in run
for tr, logit in self._traces(*args, **kwargs):
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/mcmc.py", line 51, in _traces
for trace in self._gen_samples(self.num_samples, trace):
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/mcmc.py", line 35, in _gen_samples
trace = self.kernel.sample(trace)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 276, in sample
direction, tree_depth, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 159, in _build_tree
direction, tree_depth-1, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 178, in _build_tree
direction, tree_depth-1, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 178, in _build_tree
direction, tree_depth-1, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 155, in _build_tree
return self._build_basetree(z, r, z_grads, log_slice, direction, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 130, in _build_basetree
z, r, self._potential_energy, self._inverse_mass_matrix, step_size, z_grads=z_grads)
File "/home/travis/build/uber/pyro/pyro/ops/integrator.py", line 68, in single_step_velocity_verlet
z_grads, potential_energy = _potential_grad(potential_fn, z_next)
File "/home/travis/build/uber/pyro/pyro/ops/integrator.py", line 79, in _potential_grad
potential_energy = potential_fn(z)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/hmc.py", line 170, in _potential_energy
potential_energy = -self._compute_trace_log_prob(trace)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/hmc.py", line 154, in _compute_trace_log_prob
return self._trace_prob_evaluator.log_prob(model_trace)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/util.py", line 132, in log_prob
return model_trace.log_prob_sum()
File "/home/travis/build/uber/pyro/pyro/poutine/trace_struct.py", line 130, in log_prob_sum
log_p = site["fn"].log_prob(site["value"], *site["args"], **site["kwargs"])
File "/home/travis/virtualenv/python3.5.6/lib/python3.5/site-packages/torch/distributions/transformed_distribution.py", line 86, in log_prob
log_prob += _sum_rightmost(self.base_dist.log_prob(y),
File "/home/travis/virtualenv/python3.5.6/lib/python3.5/site-packages/torch/distributions/exponential.py", line 51, in log_prob
self._validate_sample(value)
File "/home/travis/virtualenv/python3.5.6/lib/python3.5/site-packages/torch/distributions/distribution.py", line 221, in _validate_sample
raise ValueError('The value argument must be within the support')

# Fixed code
def main(args):
    pyro.set_rng_seed(args.rng_seed)
    baseball_dataset = pd.read_csv(DATA_URL, "\t")
    train, _, player_names = train_test_split(baseball_dataset)
    at_bats, hits = train[:, 0], train[:, 1]
    nuts_kernel = NUTS(conditioned_model)
    logging.info("Original Dataset:")
    logging.info(baseball_dataset)

    # (1) Full Pooling Model
    posterior_fully_pooled = MCMC(
        nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
    ).run(fully_pooled, at_bats, hits)
    logging.info("\nModel: Fully Pooled")
    logging.info("===================")
    logging.info("\nphi:")
    logging.info(
        summary(posterior_fully_pooled, sites=["phi"], player_names=player_names)["phi"]
    )
    posterior_predictive = TracePredictive(
        fully_pooled, posterior_fully_pooled, num_samples=args.num_samples
    )
    sample_posterior_predictive(posterior_predictive, baseball_dataset)
    evaluate_log_predictive_density(
        fully_pooled, posterior_fully_pooled, baseball_dataset
    )

    # (2) No Pooling Model
    posterior_not_pooled = MCMC(
        nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
    ).run(not_pooled, at_bats, hits)
    logging.info("\nModel: Not Pooled")
    logging.info("=================")
    logging.info("\nphi:")
    logging.info(
        summary(posterior_not_pooled, sites=["phi"], player_names=player_names)["phi"]
    )
    posterior_predictive = TracePredictive(
        not_pooled, posterior_not_pooled, num_samples=args.num_samples
    )
    sample_posterior_predictive(posterior_predictive, baseball_dataset)
    evaluate_log_predictive_density(not_pooled, posterior_not_pooled, baseball_dataset)

    # (3) Partially Pooled Model
    # TODO: remove once htps://github.com/uber/pyro/issues/1458 is resolved
    if "CI" not in os.environ:
        posterior_partially_pooled = MCMC(
            nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
        ).run(partially_pooled, at_bats, hits)
        logging.info("\nModel: Partially Pooled")
        logging.info("=======================")
        logging.info("\nphi:")
        logging.info(
            summary(
                posterior_partially_pooled, sites=["phi"], player_names=player_names
            )["phi"]
        )
        posterior_predictive = TracePredictive(
            partially_pooled, posterior_partially_pooled, num_samples=args.num_samples
        )
        sample_posterior_predictive(posterior_predictive, baseball_dataset)
        evaluate_log_predictive_density(
            partially_pooled, posterior_partially_pooled, baseball_dataset
        )

    # (4) Partially Pooled with Logit Model
    posterior_partially_pooled_with_logit = MCMC(
        nuts_kernel, num_samples=args.num_samples, warmup_steps=args.warmup_steps
    ).run(partially_pooled_with_logit, at_bats, hits)
    logging.info("\nModel: Partially Pooled with Logit")
    logging.info("==================================")
    logging.info("\nSigmoid(alpha):")
    logging.info(
        summary(
            posterior_partially_pooled_with_logit,
            sites=["alpha"],
            player_names=player_names,
            transforms={"alpha": lambda x: 1.0 / (1 + np.exp(-x))},
        )["alpha"]
    )
    posterior_predictive = TracePredictive(
        partially_pooled_with_logit,
        posterior_partially_pooled_with_logit,
        num_samples=args.num_samples,
    )
    sample_posterior_predictive(posterior_predictive, baseball_dataset)
    evaluate_log_predictive_density(
        partially_pooled_with_logit,
        posterior_partially_pooled_with_logit,
        baseball_dataset,
    )