def test_sorted_piecewise_constant_pdf_large_flat(self):
    """Test sampling when given a large flat distribution."""
    num_samples = 100
    num_bins = 100000
    key = random.PRNGKey(0)
    bins = jnp.arange(num_bins)
    weights = np.ones(len(bins) - 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)

    # All samples should collectively resemble a uniform distribution.
    self.assertLessEqual(
        sp.stats.kstest(samples, 'uniform', (bins[0], bins[-1])).statistic, 0.2)
Test sampling when given a large flat distribution.
internal/math_test.py