code: def test_log_prob_unit_rate(
      num_leaves, num_steps, batch_shape, sample_shape
  ):
      leaf_times = (
          torch.rand(batch_shape + (num_leaves,)).pow(0.5) * num_steps
      )
      d1 = CoalescentTimes(leaf_times)
  
      rate_grid = torch.ones(batch_shape + (num_steps,))
      d2 = CoalescentTimesWithRate(leaf_times, rate_grid)
  
      coal_times = d1.sample(sample_shape)
      assert_close(d1.log_prob(coal_times), d2.log_prob(coal_times))
  
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 4 for calling a global function with example usage: test_log_prob_unit_rate(num_leaves, num_steps, batch_shape, sample_shape) without return types