def testSoftmaxWhereMask(self, fn):
    x = jnp.array([5.5, 1.3, -4.2, 0.9])
    m = jnp.array([True, False, True, True])
    x_filtered = jnp.take(x, jnp.array([0, 2, 3]))

    out_masked = jnp.take(
        fn(x, where=m, initial=-jnp.inf), jnp.array([0, 2, 3]))
    out_filtered = fn(x_filtered)

    self.assertAllClose(out_masked, out_filtered)