ENH: stats.epps_singleton_2samp: add array API support

[mdhaber]

Reference issue

Toward gh-20544

What does this implement/fix?

Adds array API support to scipy.stats.epps_singleton_2samp.

Additional information

Updates array_api_extra to bring in the new vectorized cov.

[mdhaber]

Most CI failures are data-apis/array-api-extra#454. I'll wait for that to be resolved before re-running CI.

[lucascolley]

The following rudimentary diff makes the tests pass with NumPy for me:

diff --git a/scipy/spatial/transform/_rigid_transform.py b/scipy/spatial/transform/_rigid_transform.py
index 4d40b98ea4..e58954fee4 100644
--- a/scipy/spatial/transform/_rigid_transform.py
+++ b/scipy/spatial/transform/_rigid_transform.py
@@ -410,3 +410,3 @@ class RigidTransform:
         if self._single:
-            matrix = xpx.atleast_nd(matrix, ndim=3, xp=xp)
+            matrix = xp.expand_dims(matrix, axis=0)

@@ -993,5 +993,14 @@ class RigidTransform:
         xp = array_namespace(transforms[0].as_matrix())
-        matrix = xp.concat(
-            [xpx.atleast_nd(x.as_matrix(), ndim=3, xp=xp) for x in transforms]
-        )
+        expanded_transforms = []
+        for x in transforms:
+            x = x.as_matrix()
+            if x.ndim == 3:
+                expanded_transforms.append(x)
+            elif x.ndim == 2:
+                expanded_transforms.append(xp.expand_dims(x, axis=0))
+            else:
+                raise ValueError("panic")
+        matrix = xp.concat(expanded_transforms)
         return RigidTransform._from_raw_matrix(matrix, xp, None)
@@ -1936,3 +1945,3 @@ class RigidTransform:
         if tf._single:
-            matrix = xpx.atleast_nd(matrix, ndim=3, xp=xp)
+            matrix = xp.expand_dims(matrix, axis=0)
         tf._matrix = matrix

@scottshambaugh @amacati what do you think about merging a polished version of this diff?

Long-term I think we should just enforce that atleast_nd adds dimensions to the front, so the existing code would be fine. But right now it is relying on undefined behaviour.

EDIT: context: data-apis/array-api-extra#454 (comment)

[amacati]

At first glance this seems incorrect.

+            x = x.as_matrix()
+            if x.ndim == 3:
+                expanded_transforms.append(x)
+            elif x.ndim == 2:
+                expanded_transforms.append(xp.expand_dims(x, axis=0))
+            else:
+                raise ValueError("panic")

RigidTransform can be a batch of multiple leading dimensions. Concatenating several 4D RigidTransforms is well-defined, but would error with this loop, right?

[lucascolley]

in which case we should add test coverage for that too!

[amacati]

Yes, I am surprised that this was not picked up already.

[amacati]

Yeah, looking at the concatenate tests, the Nd case is not covered. It would also make sense to introduce an axis argument, but that's maybe something for later, or should be implemented if people actually need it. I will try to expand the test coverage asap, just a bit short of time these days.

[mdhaber]

@amacati @lucascolley I merged main here and expected the spatial tests to pass based on the discussion and linked PRs, but I still see failures. In the PR to fix them, could you update array-api-extra to make sure tests still pass?

[mschuck-rai]

Huh, since when are these tests failing? Does this come with the latest version of array-api-extra that fixes atleast_nd?

[mdhaber]

since when are these tests failing

Since I tried to update array-api-extra in this PR. That's what motivated the discussion above (#23929 (comment)).

[mdhaber]

I didn't realize array-api-extra was updated, too; I thought it was just #23937. I'll update that in this PR.

[lucascolley]

yeah, np.atleast_3d has some quite odd semantics, simply avoiding delegating to NumPy seemed like the most reasonable way forward.

[lucascolley]

LGTM, thanks Matt. I might look at doing an array-api-extra release, it has been a while and there is something nice about bumping to a tagged commit.