def test_sorted_piecewise_constant_pdf_single_bin(self):
    """Test sampling when given a small `one hot' distribution."""
    num_samples = 625
    key = random.PRNGKey(0)
    bins = jnp.array([0, 1, 3, 6, 10], jnp.float32)
    for randomized in [False, True]:
      for i in range(len(bins) - 1):
        weights = np.zeros(len(bins) - 1, jnp.float32)
        weights[i] = 1.
        samples = math.sorted_piecewise_constant_pdf(
            key,
            bins[None],
            weights[None],
            num_samples,
            randomized,
        )[0]

        # All samples should be within [bins[i], bins[i+1]].
        self.assertTrue(jnp.all(samples >= bins[i]))
        self.assertTrue(jnp.all(samples <= bins[i + 1]))
Test sampling when given a small `one hot' distribution.
internal/math_test.py