Repository: astropy/astropy

Issue:
# BUG: Fix broadcasting bug in uniform() for ndim >= 2 inputs


<!-- These comments are hidden when you submit the pull request,
so you do not need to remove them! -->

<!-- Please be sure to check out our contributing guidelines,
https://github.com/astropy/astropy/blob/main/CONTRIBUTING.md .
Please be sure to check out our code of conduct,
https://github.com/astropy/astropy/blob/main/CODE_OF_CONDUCT.md . -->

<!-- If you are new or need to be re-acquainted with Astropy
contributing workflow, please see
https://docs.astropy.org/en/latest/development/quickstart.html .
There is even a practical example at
https://docs.astropy.org/en/latest/development/git_edit_workflow_examples.html . -->

<!-- Please just have a quick search on GitHub to see if a similar
pull request has already been posted.
We have old closed pull requests that might provide useful code or ideas
that directly tie in with your pull request. -->

<!-- We have several automatic features that run when a pull request is open.
They can appear daunting but do not worry because maintainers will help
you navigate them, if necessary. -->

### Description
<!-- Provide a general description of what your pull request does.
Complete the following sentence and add relevant details as you see fit. -->

<!-- In addition please ensure that the pull request title is descriptive
and allows maintainers to infer the applicable subpackage(s). -->

<!-- READ THIS FOR MANUAL BACKPORT FROM A MAINTAINER:
Apply "skip-basebranch-check" label **before** you open the PR! -->

This PR fixes a broadcasting bug in uniform() for inputs with ndim >= 2.

The sampling axis was incorrectly inserted using [:, np.newaxis], which fails for multi-dimensional arrays. It is now appended using [..., np.newaxis], consistent with normal() and poisson(). A regression test is included

<!-- If the pull request closes any open issues you can add this.
If you replace <Issue Number> with a number, GitHub will automatically link it.
If this pull request is unrelated to any issues, please remove
the following line. -->


<!-- Optional opt-out -->

- [ ] By checking this box, the PR author has requested that maintainers do **NOT** use the "Squash and Merge" button. Maintainers should respect this when possible; however, the final decision is at the discretion of the maintainer that merges the PR.


Implementation patch:
diff --git a/astropy/uncertainty/distributions.py b/astropy/uncertainty/distributions.py
index 91f617e78c6c..d4687b95bf91 100644
--- a/astropy/uncertainty/distributions.py
+++ b/astropy/uncertainty/distributions.py
@@ -199,11 +199,8 @@ def uniform(
         )
 
     newshape = lower.shape + (n_samples,)
-    if lower.shape == () and upper.shape == ():
-        width = upper - lower  # scalar
-    else:
-        width = (upper - lower)[:, np.newaxis]
-        lower = lower[:, np.newaxis]
+    width = (upper - lower)[..., np.newaxis]
+    lower = lower[..., np.newaxis]
     samples = lower + width * np.random.uniform(size=newshape)
 
     return cls(samples, **kwargs)


Test patch:
diff --git a/astropy/uncertainty/tests/test_distribution.py b/astropy/uncertainty/tests/test_distribution.py
index 9fe47457e569..8d2b1de4f0d3 100644
--- a/astropy/uncertainty/tests/test_distribution.py
+++ b/astropy/uncertainty/tests/test_distribution.py
@@ -688,3 +688,29 @@ def setup_class(cls):
         cls.d = cls.d * u.Unit("km,m")
         cls.item = cls.item * u.Unit("Mm,km")
         cls.b_item = cls.b_item * u.km
+
+
+def test_helper_uniform_multidim_regression():
+    """
+    Regression test: uniform() must support ndim >= 2 inputs.
+
+    Before fix (using [:, np.newaxis]) this raised a broadcasting ValueError.
+    After fix (using [..., np.newaxis]) it should work and preserve trailing
+    sample axis.
+    """
+    lower = np.zeros((3, 4))
+    upper = np.ones((3, 4))
+    n_samples = 7
+
+    d = ds.uniform(lower=lower, upper=upper, n_samples=n_samples)
+
+    # Public shape hides sample axis
+    assert d.shape == (3, 4)
+    assert d.n_samples == n_samples
+
+    # Underlying samples must include trailing sampling axis
+    assert d.distribution.shape == (3, 4, n_samples)
+
+    # Bounds must hold for all samples
+    assert np.all(d.distribution >= lower[..., np.newaxis])
+    assert np.all(d.distribution <= upper[..., np.newaxis])