code: def test_with_rate_smoke(
      num_leaves,
      num_steps,
      leaf_times_shape,
      rate_grid_shape,
      sample_shape,
  ):
      batch_shape = broadcast_shape(leaf_times_shape, rate_grid_shape)
      leaf_times = (
          torch.rand(leaf_times_shape + (num_leaves,)).pow(0.5)
          * num_steps
      )
      rate_grid = torch.rand(rate_grid_shape + (num_steps,))
      d = CoalescentTimesWithRate(leaf_times, rate_grid)
      coal_times = _sample_coalescent_times(
          leaf_times.expand(sample_shape + batch_shape + (-1,))
      )
      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 3 for calling a global function with example usage: test_with_rate_smoke(num_leaves, num_steps, leaf_times_shape, rate_grid_shape, sample_shape) without return types