code: def test_log_prob_constant_rate_2(
      num_leaves, num_steps, batch_shape, sample_shape
  ):
      rate = torch.randn(batch_shape).exp()
      rate_grid = rate.unsqueeze(-1).expand(batch_shape + (num_steps,))
      leaf_times = (
          torch.rand(batch_shape + (num_leaves,)).pow(0.5) * num_steps
      )
  
      d1 = CoalescentTimes(leaf_times, rate)
      coal_times = d1.sample(sample_shape)
      log_prob_1 = d1.log_prob(coal_times)
  
      d2 = CoalescentTimesWithRate(leaf_times, rate_grid)
      log_prob_2 = d2.log_prob(coal_times)
  
      assert_close(log_prob_1, log_prob_2)
  
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 7 for calling a global function with example usage: test_log_prob_constant_rate_2(num_leaves, num_steps, batch_shape, sample_shape) without return types