ENH: Adding dtype parameter to the test for linalg.trace

[prady0t]

Found in #439 (Please also refer to it for all missing coverage reports)

This PR, along with #442, adds tests for tril and triu. Originally, only their signature tests were available. Also, in the above findings, linalg.trace wasn't tested on non-default arguments of dtype (only None); this PR addresses that.

[ev-br]

Nice!

Let's split the PR into two though: one adding tril/triu tests, and a separate one for trace. Both need some work, too, per array-api-compat testing below.

[ev-br]

Running

$ ARRAY_API_TESTS_SKIP_DTYPES=uint16,uint32,uint64 ARRAY_API_TESTS_MODULE=array_api_compat.torch pytest array_api_tests/test_linalg.py::test_trace  -sv --max-examples=10_000

emits

array_api_tests/test_linalg.py::test_trace
  /home/br/repos/array-api-compat/src/array_api_compat/torch/_aliases.py:360: UserWarning: Casting complex values to real discards the imaginary part (Triggered internally at /pytorch/aten/src/ATen/native/Copy.cpp:308.)
    return torch.sum(x, axis, dtype=dtype, keepdims=keepdims, **kwargs)

With numpy, a similar stanza,

$ ARRAY_API_TESTS_SKIP_DTYPES=uint16,uint32,uint64 ARRAY_API_TESTS_MODULE=array_api_compat.numpy pytest array_api_tests/test_linalg.py::test_trace  -sv --max-examples=10_000

generates both a ComplexWarning, and a failure:

array_api_tests/test_linalg.py::test_trace FAILED

======================================= FAILURES =======================================
______________________________________ test_trace ______________________________________

    @pytest.mark.unvectorized
>   @pytest.mark.xp_extension('linalg')

array_api_tests/test_linalg.py:918: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

x = array([[0]], dtype=uint8), kw = {}

    @pytest.mark.unvectorized
    @pytest.mark.xp_extension('linalg')
    @given(
        x=arrays(dtype=numeric_dtypes, shape=matrix_shapes()),
        # offset may produce an overflow if it is too large. Supporting offsets
        # that are way larger than the array shape isn't very important.
        kw=kwargs(offset=integers(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE),
                  dtype=sampled_from(dh.numeric_dtypes))
    )
    def test_trace(x, kw):
        res = linalg.trace(x, **kw)
    
        dtype = kw.get("dtype", None)
        expected_dtype = dh.accumulation_result_dtype(x.dtype, dtype)
        if expected_dtype is None:
            # If a default uint cannot exist (i.e. in PyTorch which doesn't support
            # uint32 or uint64), we skip testing the output dtype.
            # See https://github.com/data-apis/array-api-tests/issues/160
            if x.dtype in dh.uint_dtypes:
>               assert dh.is_int_dtype(res.dtype)  # sanity check
E               AssertionError: assert False
E                +  where False = <function is_int_dtype at 0x7f85b7900e00>(dtype('uint64'))
E                +    where <function is_int_dtype at 0x7f85b7900e00> = dh.is_int_dtype
E                +    and   dtype('uint64') = array(0, dtype=uint64).dtype
E               Falsifying example: test_trace(
E                   x=_f(
E                       <...>,  # or any other generated value
E                       False,  # or any other generated value
E                   ),
E                   kw={},
E               )
E               Explanation:
E                   These lines were always and only run by failing examples:
E                       /home/br/repos/array-api-tests/array_api_tests/dtype_helpers.py:323

array_api_tests/test_linalg.py:936: AssertionError

[prady0t]

This fails on the master branch too. So it should be unrelated to the current changes. I believe it fails because the only allowed x.dtype is uint8, whereas res.dtype could be anything, in this case, uint64. Trying to print some values before the error:

kw=  {}
x.dtype = uint8
dtype = None
res.dtype = uint64
expected_dtype = None
int_dtypes = (<class 'numpy.int8'>, <class 'numpy.int16'>, <class 'numpy.int32'>, <class 'numpy.int64'>)
uint_dtypes = (<class 'numpy.uint8'>,)
available_dtypes = (<class 'numpy.uint8'>, <class 'numpy.int8'>, <class 'numpy.int16'>, <class 'numpy.int32'>, <class 'numpy.int64'>, <class 'numpy.float32'>, <class 'numpy.float64'>, <class 'numpy.complex64'>, <class 'numpy.complex128'>)

[prady0t]

This fails on the master branch too
@ev-br

[ev-br]

Okay, checked locally: the error is indeed a red herring, is due to skipping unsigned + numpy returning a uint64 (which is perfectly fine in itself). Resolving

[ev-br]

LGTM now, thanks @prady0t !