def test_sorted_piecewise_constant_pdf_sparse_delta(self):
    """Test sampling when given a large distribution with a big delta in it."""
    num_samples = 100
    num_bins = 100000
    key = random.PRNGKey(0)
    bins = jnp.arange(num_bins)
    weights = np.ones(len(bins) - 1)
    delta_idx = len(weights) // 2
    weights[delta_idx] = len(weights) - 1
    samples = math.sorted_piecewise_constant_pdf(
        key,
        bins[None],
        weights[None],
        num_samples,
        True,
    )[0]

    # All samples should be within the range of the bins.
    self.assertTrue(jnp.all(samples >= bins[0]))
    self.assertTrue(jnp.all(samples <= bins[-1]))

    # Samples modded by their bin index should resemble a uniform distribution.
    samples_mod = jnp.mod(samples, 1)
    self.assertLessEqual(
        sp.stats.kstest(samples_mod, 'uniform', (0, 1)).statistic, 0.2)

    # The delta function bin should contain ~half of the samples.
    in_delta = (samples >= bins[delta_idx]) & (samples <= bins[delta_idx + 1])
    self.assertAllClose(jnp.mean(in_delta), 0.5, atol=0.05)
Test sampling when given a large distribution with a big delta in it.
internal/math_test.py