def test_sorted_piecewise_constant_pdf_train_mode(self):
    """Test that piecewise-constant sampling reproduces its distribution."""
    batch_size = 4
    num_bins = 16
    num_samples = 1000000
    precision = 1e5
    rng = random.PRNGKey(20202020)

    # Generate a series of random PDFs to sample from.
    data = []
    for _ in range(batch_size):
      rng, key = random.split(rng)
      # Randomly initialize the distances between bins.
      # We're rolling our own fixed precision here to make cumsum exact.
      bins_delta = jnp.round(precision * jnp.exp(
          random.uniform(key, shape=(num_bins + 1,), minval=-3, maxval=3)))

      # Set some of the bin distances to 0.
      rng, key = random.split(rng)
      bins_delta *= random.uniform(key, shape=bins_delta.shape) < 0.9

      # Integrate the bins.
      bins = jnp.cumsum(bins_delta) / precision
      rng, key = random.split(rng)
      bins += random.normal(key) * num_bins / 2
      rng, key = random.split(rng)

      # Randomly generate weights, allowing some to be zero.
      weights = jnp.maximum(
          0, random.uniform(key, shape=(num_bins,), minval=-0.5, maxval=1.))
      gt_hist = weights / weights.sum()
      data.append((bins, weights, gt_hist))

    # Tack on an "all zeros" weight matrix, which is a common cause of NaNs.
    weights = jnp.zeros_like(weights)
    gt_hist = jnp.ones_like(gt_hist) / num_bins
    data.append((bins, weights, gt_hist))

    bins, weights, gt_hist = [jnp.stack(x) for x in zip(*data)]

    for randomized in [True, False]:
      rng, key = random.split(rng)
      # Draw samples from the batch of PDFs.
      samples = math.sorted_piecewise_constant_pdf(
          key,
          bins,
          weights,
          num_samples,
          randomized,
      )
      self.assertEqual(samples.shape[-1], num_samples)

      # Check that samples are sorted.
      self.assertTrue(jnp.all(samples[..., 1:] >= samples[..., :-1]))

      # Verify that each set of samples resembles the target distribution.
      for i_samples, i_bins, i_gt_hist in zip(samples, bins, gt_hist):
        i_hist = jnp.float32(jnp.histogram(i_samples, i_bins)[0]) / num_samples
        i_gt_hist = jnp.array(i_gt_hist)

        # Merge any of the zero-span bins until there aren't any left.
        while jnp.any(i_bins[:-1] == i_bins[1:]):
          j = int(jnp.where(i_bins[:-1] == i_bins[1:])[0][0])
          i_hist = jnp.concatenate([
              i_hist[:j],
              jnp.array([i_hist[j] + i_hist[j + 1]]), i_hist[j + 2:]
          ])
          i_gt_hist = jnp.concatenate([
              i_gt_hist[:j],
              jnp.array([i_gt_hist[j] + i_gt_hist[j + 1]]), i_gt_hist[j + 2:]
          ])
          i_bins = jnp.concatenate([i_bins[:j], i_bins[j + 1:]])

        # Angle between the two histograms in degrees.
        angle = 180 / jnp.pi * jnp.arccos(
            jnp.minimum(
                1.,
                jnp.mean(
                    (i_hist * i_gt_hist) /
                    jnp.sqrt(jnp.mean(i_hist**2) * jnp.mean(i_gt_hist**2)))))
        # Jensen-Shannon divergence.
        m = (i_hist + i_gt_hist) / 2
        js_div = jnp.sum(
            sp.special.kl_div(i_hist, m) + sp.special.kl_div(i_gt_hist, m)) / 2
        self.assertLessEqual(angle, 0.5)
        self.assertLessEqual(js_div, 1e-5)
Test that piecewise-constant sampling reproduces its distribution.
internal/math_test.py