[galleonli]

Types of changes

Bug fix.

Motivation and context / Related issue

Related to #518.

wasserstein_1d accepts values with shape (n, ...), but passing more than one batch axis currently fails in quantile_function. Returning plans also fails for ordinary 1D inputs and for batches passed with require_sort=False.

This flattens the batch axes for searchsorted, then restores the indices to the input layout before gathering the quantiles. Weights shared across omitted batch axes are expanded along each missing axis. Plan construction uses one column per batch and only maps indices through a sorter when sorting was needed.

The existing list of plans is kept. Multiple batch axes are flattened in row-major order; an unbatched input returns a list containing one plan. Plan rows and columns still refer to the original sample order. The generic backend searchsorted methods are unchanged, so this does not address every part of #518.

How has this been tested (if it applies)

• Added 128 parametrized cases across NumPy, PyTorch and JAX, covering unbatched inputs, multiple batch axes, shared and per-batch weights, zero weights, pre-sorted inputs, and both plan return formats. Against the original solver, 104 cases fail and 24 pass.
• Compared transport plans, marginal sums and costs with ot.emd; checked partially shared weights against ot.emd2 for both p=1 and p=2.
• Compared PyTorch and JAX gradients with independent 1D calls. PyTorch gradcheck and the JAX JIT comparison pass, including non-contiguous PyTorch inputs.
• The 1D solver, sliced modules and related unbalanced tests pass: 328 passed, 39 skipped. The separate circle solver module has 18 passed.
• All five applicable pre-commit hooks and git diff --check pass.

These were local CPU runs with POT's native extensions built from the checked-out source. The full test suite, documentation build, GPU execution, TensorFlow and CuPy were not tested.

PR checklist

• I have read the CONTRIBUTING document.
• The documentation is up-to-date with the changes I made (check build artifacts).
• All tests passed, and additional code has been covered with new tests.
• I have added the PR and Issue fix to the RELEASES.md file.

The docstrings and release note are updated. The two unchecked items reflect the documentation build and full test suite not being run locally; the targeted results are listed above.