Merge pull request #434 from ev-br/torch_mps

ENH: allow running on torch.mps device, which does not have float64/complex128

Author: ev-br

array_api_tests/hypothesis_helpers.py

@@ -341,8 +341,8 @@ def finite_matrices(draw, shape=matrix_shapes(), dtype=floating_dtypes, bound=No
 
 
 rtol_shared_matrix_shapes = shared(matrix_shapes())
-# Should we set a max_value here?
-_rtol_float_kw = dict(allow_nan=False, allow_infinity=False, min_value=0)
+# Arbitrary max_value for rtols, to avoid overflows when float64 is not available
+_rtol_float_kw = dict(allow_nan=False, allow_infinity=False, min_value=0, max_value=42)
 rtols = one_of(floats(**_rtol_float_kw),
                arrays(dtype=real_floating_dtypes,
                           shape=rtol_shared_matrix_shapes.map(lambda shape:  shape[:-2]),

array_api_tests/test_signatures.py

@@ -156,7 +156,10 @@ def make_pretty_func(func_name: str, *args: Any, **kwargs: Any) -> str:
     array_argnames -= set(func_to_specified_arg_exprs[func_name].keys())
     if len(array_argnames) > 0:
         in_dtypes = dh.func_in_dtypes[func_name]
-        for dtype_name in ["float64", "bool", "int64", "complex128"]:
+        # use "float64" if available, "float32" otherwise; ditto for complex128/complex64
+        float_name = dh.dtype_to_name[dh.widest_real_dtype]
+        cmplx_name = dh.dtype_to_name[dh.widest_complex_dtype]
+        for dtype_name in [float_name, "bool", "int64", cmplx_name]:
             # We try float64 first because uninspectable numerical functions
             # tend to support float inputs first-and-foremost (i.e. PyTorch)
             try: