Repository: data-apis/array-api-strict

Issue:
`squeeze` accepts `axis=None`?
It looks like `axis=None` is not in the standard for [`squeeze`](https://data-apis.org/array-api/2023.12/API_specification/generated/array_api.squeeze.html#array_api.squeeze), yet it is accepted by `array_api_strict`.
```python3
import array_api_strict as xp
x = xp.reshape(xp.arange(12), (3, 4))[..., xp.newaxis]
xp.squeeze(x, axis=None)
# Array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11]], dtype=array_api_strict.int64)
```

Discussion:
@ev-br: I suspect it's an omission in the standard. ISTM all array libraries allow `axis=None` with the meaning of "remove all singleton dimensions":

```
In [17]: t = torch.empty((1, 2, 1, 3, 4, 1))

In [18]: torch.squeeze(t).shape
Out[18]: torch.Size([2, 3, 4])
```

torch does not allow `squeeze(t, dim=None)` which is fairly typical for torch's implementations of optional arguments, so will need to be worked around in the `array-api-compat` layer.

cross-ref https://github.com/data-apis/array-api/issues/892@ev-br: Well, no, apparently, it's not an omission in the standard, it's deliberate, cf https://github.com/data-apis/array-api/issues/892#issuecomment-2634685099

Let's then block axis=None in `-strict` : https://github.com/data-apis/array-api-strict/pull/125

Solution patch:
diff --git a/array_api_strict/_array_object.py b/array_api_strict/_array_object.py
index a917441..afee030 100644
--- a/array_api_strict/_array_object.py
+++ b/array_api_strict/_array_object.py
@@ -26,10 +26,12 @@ from ._dtypes import (
     _integer_dtypes,
     _integer_or_boolean_dtypes,
     _floating_dtypes,
+    _real_floating_dtypes,
     _complex_floating_dtypes,
     _numeric_dtypes,
     _result_type,
     _dtype_categories,
+    _real_to_complex_map,
 )
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
 
@@ -243,6 +245,7 @@ class Array:
         """
         from ._data_type_functions import iinfo
 
+        target_dtype = self.dtype
         # Note: Only Python scalar types that match the array dtype are
         # allowed.
         if isinstance(scalar, bool):
@@ -268,10 +271,13 @@ class Array:
                     "Python float scalars can only be promoted with floating-point arrays."
                 )
         elif isinstance(scalar, complex):
-            if self.dtype not in _complex_floating_dtypes:
+            if self.dtype not in _floating_dtypes:
                 raise TypeError(
-                    "Python complex scalars can only be promoted with complex floating-point arrays."
+                    "Python complex scalars can only be promoted with floating-point arrays."
                 )
+            # 1j * array(floating) is allowed
+            if self.dtype in _real_floating_dtypes:
+                target_dtype = _real_to_complex_map[self.dtype]
         else:
             raise TypeError("'scalar' must be a Python scalar")
 
@@ -282,7 +288,7 @@ class Array:
         # behavior for integers within the bounds of the integer dtype.
         # Outside of those bounds we use the default NumPy behavior (either
         # cast or raise OverflowError).
-        return Array._new(np.array(scalar, dtype=self.dtype._np_dtype), device=self.device)
+        return Array._new(np.array(scalar, dtype=target_dtype._np_dtype), device=self.device)
 
     @staticmethod
     def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
diff --git a/array_api_strict/_data_type_functions.py b/array_api_strict/_data_type_functions.py
index 5af46d2..1643043 100644
--- a/array_api_strict/_data_type_functions.py
+++ b/array_api_strict/_data_type_functions.py
@@ -197,7 +197,7 @@ def isdtype(
     else:
         raise TypeError(f"'kind' must be a dtype, str, or tuple of dtypes and strs, not {type(kind).__name__}")
 
-def result_type(*arrays_and_dtypes: Union[Array, Dtype]) -> Dtype:
+def result_type(*arrays_and_dtypes: Union[Array, Dtype, int, float, complex, bool]) -> Dtype:
     """
     Array API compatible wrapper for :py:func:`np.result_type <numpy.result_type>`.
 
@@ -208,19 +208,40 @@ def result_type(*arrays_and_dtypes: Union[Array, Dtype]) -> Dtype:
     # too many extra type promotions like int64 + uint64 -> float64, and does
     # value-based casting on scalar arrays.
     A = []
+    scalars = []
     for a in arrays_and_dtypes:
         if isinstance(a, Array):
             a = a.dtype
+        elif isinstance(a, (bool, int, float, complex)):
+            scalars.append(a)
         elif isinstance(a, np.ndarray) or a not in _all_dtypes:
             raise TypeError("result_type() inputs must be array_api arrays or dtypes")
         A.append(a)
 
+    # remove python scalars
+    A = [a for a in A if not isinstance(a, (bool, int, float, complex))]
+
     if len(A) == 0:
         raise ValueError("at least one array or dtype is required")
     elif len(A) == 1:
-        return A[0]
+        result = A[0]
     else:
         t = A[0]
         for t2 in A[1:]:
             t = _result_type(t, t2)
-        return t
+        result = t
+
+    if len(scalars) == 0:
+        return result
+
+    if get_array_api_strict_flags()['api_version'] <= '2023.12':
+        raise TypeError("result_type() inputs must be array_api arrays or dtypes")
+
+    # promote python scalars given the result_type for all arrays/dtypes
+    from ._creation_functions import empty
+    arr = empty(1, dtype=result)
+    for s in scalars:
+        x = arr._promote_scalar(s)
+        result = _result_type(x.dtype, result)
+
+    return result
diff --git a/array_api_strict/_dtypes.py b/array_api_strict/_dtypes.py
index b51ed92..66304dd 100644
--- a/array_api_strict/_dtypes.py
+++ b/array_api_strict/_dtypes.py
@@ -126,6 +126,7 @@ _dtype_categories = {
     "floating-point": _floating_dtypes,
 }
 
+_real_to_complex_map = {float32: complex64, float64: complex128}
 
 # Note: the spec defines a restricted type promotion table compared to NumPy.
 # In particular, cross-kind promotions like integer + float or boolean +
diff --git a/array_api_strict/_fft.py b/array_api_strict/_fft.py
index 4b0ceb6..c888826 100644
--- a/array_api_strict/_fft.py
+++ b/array_api_strict/_fft.py
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Union, Optional, Literal
-    from ._typing import Device
+    from ._typing import Device, Dtype as DType
     from collections.abc import Sequence
 
 from ._dtypes import (
@@ -251,7 +251,14 @@ def ihfft(
     return res
 
 @requires_extension('fft')
-def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Array:
+def fftfreq(
+    n: int,
+    /,
+    *,
+    d: float = 1.0,
+    dtype: Optional[DType] = None,
+    device: Optional[Device] = None
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.fftfreq <numpy.fft.fftfreq>`.
 
@@ -259,10 +266,23 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Ar
     """
     if device is not None and device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
-    return Array._new(np.fft.fftfreq(n, d=d), device=device)
+    if dtype and not dtype in _real_floating_dtypes:
+        raise ValueError(f"`dtype` must be a real floating-point type. Got {dtype=}.")
+
+    np_result = np.fft.fftfreq(n, d=d)
+    if dtype:
+        np_result = np_result.astype(dtype._np_dtype)
+    return Array._new(np_result, device=device)
 
 @requires_extension('fft')
-def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Array:
+def rfftfreq(
+    n: int,
+    /,
+    *,
+    d: float = 1.0,
+    dtype: Optional[DType] = None,
+    device: Optional[Device] = None
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.rfftfreq <numpy.fft.rfftfreq>`.
 
@@ -270,7 +290,13 @@ def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> A
     """
     if device is not None and device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
-    return Array._new(np.fft.rfftfreq(n, d=d), device=device)
+    if dtype and not dtype in _real_floating_dtypes:
+        raise ValueError(f"`dtype` must be a real floating-point type. Got {dtype=}.")
+
+    np_result = np.fft.rfftfreq(n, d=d)
+    if dtype:
+        np_result = np_result.astype(dtype._np_dtype)
+    return Array._new(np_result, device=device)
 
 @requires_extension('fft')
 def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:
diff --git a/array_api_strict/_manipulation_functions.py b/array_api_strict/_manipulation_functions.py
index d775835..63c3516 100644
--- a/array_api_strict/_manipulation_functions.py
+++ b/array_api_strict/_manipulation_functions.py
@@ -153,6 +153,11 @@ def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
 
     See its docstring for more information.
     """
+    if axis is None:
+        raise ValueError(
+            "squeeze(..., axis=None is not supported. See "
+            "https://github.com/data-apis/array-api/pull/100 for a discussion."
+        )
     return Array._new(np.squeeze(x._array, axis=axis), device=x.device)
 
 


Test patch:
diff --git a/array_api_strict/tests/test_array_object.py b/array_api_strict/tests/test_array_object.py
index 4535d99..edfa073 100644
--- a/array_api_strict/tests/test_array_object.py
+++ b/array_api_strict/tests/test_array_object.py
@@ -108,6 +108,14 @@ def _check_op_array_scalar(dtypes, a, s, func, func_name, BIG_INT=BIG_INT):
     # - a Python int or float for real floating-point array dtypes
     # - a Python int, float, or complex for complex floating-point array dtypes
 
+    # an exception: complex scalar <op> floating array
+    scalar_types_for_float = [float, int]
+    if not (func_name.startswith("__i")
+            or (func_name in ["__floordiv__", "__rfloordiv__", "__mod__", "__rmod__"]
+                and type(s) == complex)
+    ):
+        scalar_types_for_float += [complex]
+
     if ((dtypes == "all"
          or dtypes == "numeric" and a.dtype in _numeric_dtypes
          or dtypes == "real numeric" and a.dtype in _real_numeric_dtypes
@@ -121,7 +129,7 @@ def _check_op_array_scalar(dtypes, a, s, func, func_name, BIG_INT=BIG_INT):
         # isinstance here.
         and (a.dtype in _boolean_dtypes and type(s) == bool
              or a.dtype in _integer_dtypes and type(s) == int
-             or a.dtype in _real_floating_dtypes and type(s) in [float, int]
+             or a.dtype in _real_floating_dtypes and type(s) in scalar_types_for_float
              or a.dtype in _complex_floating_dtypes and type(s) in [complex, float, int]
         )):
         if a.dtype in _integer_dtypes and s == BIG_INT:
diff --git a/array_api_strict/tests/test_data_type_functions.py b/array_api_strict/tests/test_data_type_functions.py
index 488eab7..863d3d4 100644
--- a/array_api_strict/tests/test_data_type_functions.py
+++ b/array_api_strict/tests/test_data_type_functions.py
@@ -6,9 +6,9 @@ from numpy.testing import assert_raises
 import numpy as np
 
 from .._creation_functions import asarray
-from .._data_type_functions import astype, can_cast, isdtype
+from .._data_type_functions import astype, can_cast, isdtype, result_type
 from .._dtypes import (
-    bool, int8, int16, uint8, float64,
+    bool, int8, int16, uint8, float64, int64
 )
 from .._flags import set_array_api_strict_flags
 
@@ -70,3 +70,22 @@ def astype_device(api_version):
     else:
         pytest.raises(TypeError, lambda: astype(a, int8, device=None))
         pytest.raises(TypeError, lambda: astype(a, int8, device=a.device))
+
+
+@pytest.mark.parametrize("api_version", ['2023.12', '2024.12'])
+def test_result_type_py_scalars(api_version):
+    if api_version <= '2023.12':
+        set_array_api_strict_flags(api_version=api_version)
+
+        with pytest.raises(TypeError):
+            result_type(int16, 3)
+    else:
+        with pytest.warns(UserWarning):
+            set_array_api_strict_flags(api_version=api_version)
+
+            assert result_type(int8, 3) == int8
+            assert result_type(uint8, 3) == uint8
+            assert result_type(float64, 3) == float64
+
+            with pytest.raises(TypeError):
+                result_type(int64, True)
diff --git a/array_api_strict/tests/test_elementwise_functions.py b/array_api_strict/tests/test_elementwise_functions.py
index 0b90f0b..cc3a2cd 100644
--- a/array_api_strict/tests/test_elementwise_functions.py
+++ b/array_api_strict/tests/test_elementwise_functions.py
@@ -233,15 +233,25 @@ def test_scalars():
         if nargs(func) != 2:
             continue
 
+        nocomplex = [
+            'atan2', 'copysign', 'floor_divide', 'hypot', 'logaddexp',  'nextafter',
+            'remainder',
+            'greater', 'less', 'greater_equal', 'less_equal', 'maximum', 'minimum',
+        ]
+
         for s in [1, 1.0, 1j, BIG_INT, False]:
             for a in _array_vals():
                 for func1 in [lambda s: func(a, s), lambda s: func(s, a)]:
-                    allowed = _check_op_array_scalar(dtypes, a, s, func1, func_name)
+
+                    if func_name in nocomplex and type(s) == complex:
+                        allowed = False
+                    else:
+                        allowed = _check_op_array_scalar(dtypes, a, s, func1, func_name)
 
                     # only check `func(array, scalar) == `func(array, array)` if
                     # the former is legal under the promotion rules
                     if allowed:
-                        conv_scalar = asarray(s, dtype=a.dtype)
+                        conv_scalar = a._promote_scalar(s)
 
                         with suppress_warnings() as sup:
                             # ignore warnings from pow(BIG_INT)
diff --git a/array_api_strict/tests/test_manipulation_functions.py b/array_api_strict/tests/test_manipulation_functions.py
index 9969651..bd247ee 100644
--- a/array_api_strict/tests/test_manipulation_functions.py
+++ b/array_api_strict/tests/test_manipulation_functions.py
@@ -11,7 +11,7 @@ from .._manipulation_functions import (
 
 
 def test_concat_errors():
-    assert_raises(TypeError, lambda: concat((1, 1), axis=None))
+    assert_raises((TypeError, ValueError), lambda: concat((1, 1), axis=None))
     assert_raises(TypeError, lambda: concat([asarray([1], dtype=int8),
                                              asarray([1], dtype=float64)]))
 
diff --git a/array_api_strict/tests/test_validation.py b/array_api_strict/tests/test_validation.py
index 035b6f4..bd76ec6 100644
--- a/array_api_strict/tests/test_validation.py
+++ b/array_api_strict/tests/test_validation.py
@@ -18,7 +18,7 @@ def p(func: Callable, *args, **kwargs):
     [
         p(xp.can_cast, 42, xp.int8),
         p(xp.can_cast, xp.int8, 42),
-        p(xp.result_type, 42),
+        p(xp.result_type, "42"),
     ],
 )
 def test_raises_on_invalid_types(func, args, kwargs):