Support multidimensional PyTorch quantiles

src/pyrecest/backend_support/_pytorch_quantile_empty_batch_contract.py

@@ -4,7 +4,7 @@
 
 
 def patch_pytorch_quantile_empty_batch_contract() -> None:
-    """Preserve empty non-reduced dimensions in PyTorch quantile reductions."""
+    """Preserve NumPy quantile shapes unsupported by native PyTorch."""
 
     try:
         import numpy as np  # pylint: disable=import-outside-toplevel
@@ -50,6 +50,7 @@ def quantile(
 
         effective_method = method if interpolation is None else interpolation
         values = raw_pytorch.array(a)
+        q_shape = raw_pytorch._quantile_q_shape(q)
         is_integral_axis = isinstance(
             effective_axis, (int, np.integer)
         ) and not isinstance(effective_axis, (bool, np.bool_))
@@ -67,7 +68,8 @@ def quantile(
                         values, dtype=raw_pytorch.get_default_dtype()
                     )
                 q_arg = raw_pytorch._quantile_q(q, values)
-                q_shape = raw_pytorch._quantile_q_shape(q)
+                if len(q_shape) > 1:
+                    q_arg = q_arg.reshape(-1)
                 validation_values = torch.zeros(
                     values.shape[normalized_axis],
                     dtype=values.dtype,
@@ -91,6 +93,34 @@ def quantile(
                     return out
                 return result
 
+        if len(q_shape) > 1:
+            quantile_values = values
+            if not raw_pytorch.is_floating(
+                quantile_values
+            ) and not raw_pytorch.is_complex(quantile_values):
+                quantile_values = raw_pytorch.cast(
+                    quantile_values,
+                    dtype=raw_pytorch.get_default_dtype(),
+                )
+            q_arg = raw_pytorch._quantile_q(q, quantile_values).reshape(-1)
+            result = original_quantile(
+                a,
+                q_arg,
+                axis=axis,
+                out=None,
+                overwrite_input=overwrite_input,
+                method=method,
+                keepdims=keepdims,
+                dim=dim,
+                keepdim=keepdim,
+                interpolation=interpolation,
+            )
+            result = result.reshape(q_shape + tuple(result.shape[1:]))
+            if out is not None:
+                out.copy_(result)
+                return out
+            return result
+
         return original_quantile(
             a,
             q,

tests/backend_support/test_pytorch_quantile_contract.py

@@ -32,6 +32,46 @@ def test_pytorch_quantile_accepts_numpy_axis_and_keepdims_keywords():
     assert result.returncode == 0, result.stderr
 
 
+@pytest.mark.backend_portable
+def test_pytorch_quantile_accepts_multidimensional_q_like_numpy():
+    if importlib.util.find_spec("torch") is None:
+        pytest.skip("PyTorch is not installed")
+
+    result = run_backend_code(
+        "pytorch",
+        """
+import numpy as np
+
+import pyrecest.backend as backend
+import pyrecest._backend.pytorch as raw_backend
+
+q = np.array([[0.25, 0.75], [0.1, 0.9]])
+integer_values = [[1, 4], [3, 8]]
+cube = np.arange(24.0).reshape(2, 3, 4)
+
+for target in (backend, raw_backend):
+    column_quantiles = target.quantile(integer_values, q, axis=0)
+    tuple_axis_quantiles = target.quantile(
+        cube,
+        q,
+        axis=(0, 2),
+        keepdims=True,
+    )
+
+    np.testing.assert_allclose(
+        backend.to_numpy(column_quantiles),
+        np.quantile(integer_values, q, axis=0),
+    )
+    np.testing.assert_allclose(
+        backend.to_numpy(tuple_axis_quantiles),
+        np.quantile(cube, q, axis=(0, 2), keepdims=True),
+    )
+""",
+    )
+
+    assert result.returncode == 0, result.stderr
+
+
 @pytest.mark.backend_portable
 def test_pytorch_quantile_preserves_empty_batch_dimensions():
     if importlib.util.find_spec("torch") is None:
@@ -51,10 +91,17 @@ def test_pytorch_quantile_preserves_empty_batch_dimensions():
     axis=1,
     keepdims=True,
 )
+matrix_quantile = backend.quantile(
+    values,
+    [[0.25, 0.75], [0.1, 0.9]],
+    axis=1,
+    keepdims=True,
+)
 raw_quantile = raw_backend.quantile(values, 0.5, dim=1, keepdim=True)
 
 assert tuple(scalar_quantile.shape) == (0, 2)
 assert tuple(vector_quantile.shape) == (2, 0, 1, 2)
+assert tuple(matrix_quantile.shape) == (2, 2, 0, 1, 2)
 assert tuple(raw_quantile.shape) == (0, 1, 2)
 """,
     )