BUG: fix support for gufuncs with more than one core dimension

astropy/utils/masked/core.py

@@ -773,7 +773,7 @@
                     ) = np.lib._function_base_impl._parse_gufunc_signature(
                         ufunc.signature.replace(" ", "")
                     )
-                axis = kwargs.get("axis", -1)
+                axis = kwargs.get("axis")
                 keepdims = kwargs.get("keepdims", False)
                 in_masks = []
                 for sig, mask in zip(in_sig, masks):
@@ -783,8 +783,13 @@
                             # value in those is masked, the output will be
                             # masked too (TODO: for multiple core dimensions
                             # this may be too strong).
+                            in_axis = (
+                                tuple(range(-1, -1 - len(sig), -1))
+                                if axis is None
+                                else axis
+                            )
                             mask = np.logical_or.reduce(
-                                mask, axis=axis, keepdims=keepdims
+                                mask, axis=in_axis, keepdims=keepdims
                             )
                         in_masks.append(mask)
 
@@ -794,7 +799,10 @@
                     if os:
                         # Output has core dimensions.  Assume all those
                         # get the same mask.
-                        result_mask = np.expand_dims(mask, axis)
+                        out_axis = (
+                            tuple(range(-1, -1 - len(os), -1)) if axis is None else axis
+                        )
+                        result_mask = np.expand_dims(mask, out_axis)
                     else:
                         result_mask = mask
                     result_masks.append(result_mask)

astropy/utils/masked/tests/test_functions.py

@@ -173,6 +173,36 @@ def test_multi_op_ufunc(self):
             assert res0.unmasked == exp[0]
             assert res0.mask == expected_mask[0]
 
+    def test_erfa_rxp(self):
+        # Regression tests for gh-16116
+        m = Masked(np.eye(3))
+        v = Masked(np.arange(6).reshape(2, 3))
+        rxp1 = erfa_ufunc.rxp(m, v)
+        exp = erfa_ufunc.rxp(m.unmasked, v.unmasked)
+        assert_array_equal(rxp1.unmasked, exp)
+        assert_array_equal(rxp1.mask, False)
+        v.mask[0, 0] = True
+        rxp2 = erfa_ufunc.rxp(m, v)
+        assert_array_equal(rxp2.unmasked, exp)
+        assert_array_equal(rxp2.mask, [[True] * 3, [False] * 3])
+        m.mask[1, 1] = True
+        v.mask[...] = False
+        rxp3 = erfa_ufunc.rxp(m, v)
+        assert_array_equal(rxp3.unmasked, exp)
+        assert_array_equal(rxp3.mask, True)
+
+    def test_erfa_rxr(self):
+        m1 = Masked(np.arange(27.0).reshape(3, 3, 3))
+        m2 = Masked(np.arange(-27.0, 0.0).reshape(3, 3, 3))
+        rxr1 = erfa_ufunc.rxr(m1, m2)
+        exp = erfa_ufunc.rxr(m1.unmasked, m2.unmasked)
+        assert_array_equal(rxr1.unmasked, exp)
+        assert_array_equal(rxr1.mask, False)
+        m1.mask[0, 1, 2] = True
+        rxr2 = erfa_ufunc.rxr(m1, m2)
+        assert_array_equal(rxr2.unmasked, exp)
+        assert np.all(rxr2.mask == [[[True]], [[False]], [[False]]])
+
     @pytest.mark.parametrize("axis", (0, 1, None))
     def test_add_reduce(self, axis):
         ma_reduce = np.add.reduce(self.ma, axis=axis)

docs/changes/utils/16120.bugfix.rst

@@ -0,0 +1,2 @@
+Fix support in ``Masked`` for generalized ufuncs with more than a
+single core dimension (such as ``erfa.rxp``).