code: def test_simple_smoke(
      num_leaves, num_steps, batch_shape, sample_shape
  ):
      leaf_times = (
          torch.rand(batch_shape + (num_leaves,)).pow(0.5) * num_steps
      )
      d = CoalescentTimes(leaf_times)
      coal_times = d.sample(sample_shape)
      assert coal_times.shape == sample_shape + batch_shape + (
          num_leaves - 1,
      )
  
      actual = d.log_prob(coal_times)
      assert actual.shape == sample_shape + batch_shape
  
imports: import io
  import re
  import pytest
  import torch
  import pyro
  from pyro.distributions import CoalescentTimes, CoalescentTimesWithRate
  from pyro.distributions.coalescent import CoalescentRateLikelihood, CoalescentTimesConstraint, _sample_coalescent_times, bio_phylo_to_times
  from pyro.distributions.util import broadcast_shape
  from tests.common import assert_close
label_desc: Use image node_id 2 for calling a global function with example usage: test_simple_smoke(num_leaves, num_steps, batch_shape, sample_shape) without return types