torch: rewrite result_type to handle scalars, add a test

Test torch-specific behavior which is unspecified in the spec.

Author: ev-br

array_api_compat/torch/_aliases.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from functools import wraps as _wraps
+from functools import reduce as _reduce, wraps as _wraps
 from builtins import all as _builtin_all, any as _builtin_any
 
 from ..common import _aliases
@@ -124,43 +124,35 @@ def _fix_promotion(x1, x2, only_scalar=True):
 
 
 def result_type(*arrays_and_dtypes: Union[array, Dtype, bool, int, float, complex]) -> Dtype:
-    if len(arrays_and_dtypes) == 0:
-        raise TypeError("At least one array or dtype must be provided")
-    if len(arrays_and_dtypes) == 1:
+    num = len(arrays_and_dtypes)
+
+    if num == 0:
+        raise ValueError("At least one array or dtype must be provided")
+
+    elif num == 1:
         x = arrays_and_dtypes[0]
         if isinstance(x, torch.dtype):
             return x
         return x.dtype
 
-    if len(arrays_and_dtypes) > 2:
-        # sort the scalars to the left so that they are treated last
-        scalars, others = [], []
-        for x in arrays_and_dtypes:
-            if isinstance(x, _py_scalars):
-                scalars.append(x)
-            else:
-                others.append(x)
-        if len(scalars) == len(arrays_and_dtypes):
-            raise ValueError("At least one array or dtype is required.")
+    if num == 2:
+        x, y = arrays_and_dtypes
+        return _result_type(x, y)
 
-        arrays_and_dtypes = scalars + others
-        return result_type(arrays_and_dtypes[0], result_type(*arrays_and_dtypes[1:]))
+    else:
+        if _builtin_all(isinstance(x, _py_scalars) for x in arrays_and_dtypes):
+            raise ValueError("At least one array or dtype must be provided")
 
-    # the binary case
-    x, y = arrays_and_dtypes
+        return _reduce(_result_type, arrays_and_dtypes)
 
-    if isinstance(x, _py_scalars):
-        if isinstance(y, _py_scalars):
-            raise ValueError("At least one array or dtype is required.")
-        return y
-    elif isinstance(y, _py_scalars):
-        return x
 
-    xdt = x.dtype if not isinstance(x, torch.dtype) else x
-    ydt = y.dtype if not isinstance(y, torch.dtype) else y
+def _result_type(x, y):
+    if not (isinstance(x, _py_scalars) or isinstance(y, _py_scalars)):
+        xdt = x.dtype if not isinstance(x, torch.dtype) else x
+        ydt = y.dtype if not isinstance(y, torch.dtype) else y
 
-    if (xdt, ydt) in _promotion_table:
-        return _promotion_table[xdt, ydt]
+        if (xdt, ydt) in _promotion_table:
+            return _promotion_table[xdt, ydt]
 
     # This doesn't result_type(dtype, dtype) for non-array API dtypes
     # because torch.result_type only accepts tensors. This does however, allow
@@ -169,6 +161,7 @@ def result_type(*arrays_and_dtypes: Union[array, Dtype, bool, int, float, comple
     y = torch.tensor([], dtype=y) if isinstance(y, torch.dtype) else y
     return torch.result_type(x, y)
 
+
 def can_cast(from_: Union[Dtype, array], to: Dtype, /) -> bool:
     if not isinstance(from_, torch.dtype):
         from_ = from_.dtype

tests/test_torch.py

@@ -0,0 +1,66 @@
+"""Test "unspecified" behavior which we cannot easily test in the Array API test suite.
+"""
+import functools
+
+import pytest
+import torch
+
+from array_api_compat import torch as xp
+
+
+class TestResultType:
+    def test_empty(self):
+        with pytest.raises(ValueError):
+            xp.result_type()
+
+    def test_one_arg(self):
+        for x in [1, 1.0, 1j, '...', None]:
+            with pytest.raises((ValueError, AttributeError)):
+                xp.result_type(x)
+
+        for x in [xp.float32, xp.int64, torch.complex64]:
+            assert xp.result_type(x) == x
+
+        for x in [xp.asarray(True, dtype=xp.bool), xp.asarray(1, dtype=xp.complex64)]:
+            assert xp.result_type(x) == x.dtype
+
+    def test_two_args(self):
+        # Only include here things "unspecified" in the spec
+
+        # scalar, tensor or tensor,tensor
+        for x, y in [
+            (1., 1j),
+            (1j, xp.arange(3)),
+            (True, xp.asarray(3.)),
+            (xp.ones(3) == 1, 1j*xp.ones(3)),
+        ]:
+            assert xp.result_type(x, y) == torch.result_type(x, y)
+
+        # dtype, scalar
+        for x, y in [
+            (1j, xp.int64),
+            (True, xp.float64),
+        ]:
+            assert xp.result_type(x, y) == torch.result_type(x, xp.empty([], dtype=y))
+
+        # dtype, dtype
+        for x, y in [
+            (xp.bool, xp.complex64)
+        ]:
+            xt, yt = xp.empty([], dtype=x), xp.empty([], dtype=y)
+            assert xp.result_type(x, y) == torch.result_type(xt, yt)
+
+    def test_multi_arg(self):
+        torch.set_default_dtype(torch.float32)
+
+        args = [1, 2, 3j, xp.arange(3), 4, 5, 6]
+        assert xp.result_type(*args) == xp.complex64
+
+        args = [1, 2, 3j, xp.float64, 4, 5, 6]
+        assert xp.result_type(*args) == xp.complex128
+
+        args = [1, 2, 3j, xp.float64, 4, xp.asarray(3, dtype=xp.int16), 5, 6, False]
+        assert xp.result_type(*args) == xp.complex128
+
+        with pytest.raises(ValueError):
+            xp.result_type(1, 2, 3, 4)