Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
prady0t wants to merge 2 commits into
base: master
Choose a base branch
from
+6 −4
Open

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

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d888d62
Adding tests for tril, triu and adding dtype parameter for linalg.trace
prady0t Jul 12, 2026
ce057d3
Moving tests for tril and triu to another PR
prady0t Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions array_api_tests/test_linalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -920,7 +920,8 @@ def test_tensordot(x1, x2, kw):
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))
kw=kwargs(offset=integers(-MAX_ARRAY_SIZE, MAX_ARRAY_SIZE),
dtype=sampled_from(dh.numeric_dtypes))

@ev-br ev-br Jul 13, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 prady0t Jul 13, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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'>)

)
def test_trace(x, kw):
res = linalg.trace(x, **kw)
Expand All @@ -939,10 +940,11 @@ def test_trace(x, kw):
n, m = x.shape[-2:]
ph.assert_result_shape('trace', x.shape, res.shape, expected=x.shape[:-2])

def true_trace(x_stack, offset=0):
def true_trace(x_stack, offset=0, dtype=None):
# Note: the spec does not specify that offset must be within the
# bounds of the matrix. A large offset should just produce a size 0
# diagonal in the last dimension (trace 0). See test_diagonal().
out_dtype = x.dtype if dtype is None else dtype
if offset < 0:
diag_size = min(n, m, max(n + offset, 0))
elif offset == 0:
Expand All @@ -954,8 +956,8 @@ def true_trace(x_stack, offset=0):
x_stack_diag = [x_stack[i, i + offset] for i in range(diag_size)]
else:
x_stack_diag = [x_stack[i - offset, i] for i in range(diag_size)]
result = xp.asarray(xp.stack(x_stack_diag) if x_stack_diag else [], dtype=x.dtype)
return _array_module.sum(result)
result = xp.asarray(xp.stack(x_stack_diag) if x_stack_diag else [], dtype=out_dtype)
return _array_module.sum(result, dtype=dtype)


_test_stacks(linalg.trace, x, **kw, res=res, dims=0, true_val=true_trace)
Expand Down
Loading