diff --git a/src/pyrecest/backend_support/_pytorch_quantile_empty_batch_contract.py b/src/pyrecest/backend_support/_pytorch_quantile_empty_batch_contract.py index d3e7a3287..4b25ed65c 100644 --- a/src/pyrecest/backend_support/_pytorch_quantile_empty_batch_contract.py +++ b/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, diff --git a/tests/backend_support/test_pytorch_quantile_contract.py b/tests/backend_support/test_pytorch_quantile_contract.py index 892954022..ea1b2b5ac 100644 --- a/tests/backend_support/test_pytorch_quantile_contract.py +++ b/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) """, )