Return a single channel from Discontinuity in 2d and 3d

[barrosyan]

Discontinuity.__call__ built its mask with jnp.ones_like(x, dtype=bool), where x is the grid of shape (D, ..., N). The mask therefore kept D channels, so in 2d and 3d RandomDiscontinuities returned D identical channels instead of the (1, ..., N) shape the other IC generators return:

import jax, exponax as ex
ex.ic.RandomDiscontinuities(2)(16, key=jax.random.PRNGKey(0)).shape
# main: (2, 16, 16)   with this change: (1, 16, 16)

This PR builds the mask from a single grid slice (x[0:1]). 1d output is unchanged.

RandomDiscontinuities was also missing from the parametrization of test_ic_output_shape, so I added it there. Its 2d and 3d cases fail on main and pass with this change, and pytest tests/test_ic.py -k Discontinuities passes (20 tests). The only usage in the docs is the 1d showcase notebook, which is not affected.

I found this while checking build_ic_set against every random IC generator for #58.