Final shape not exceed max size in hh.invertible_matrices()

honno · Mar 12, 2024 · 695c67e · 695c67e
1 parent 762f1e7
commit 695c67e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/array_api_tests/hypothesis_helpers.py b/array_api_tests/hypothesis_helpers.py
@@ -312,8 +312,8 @@ def positive_definite_matrices(draw, dtypes=xps.floating_dtypes()):
 @composite
 def invertible_matrices(draw, dtypes=xps.floating_dtypes(), stack_shapes=shapes()):
     # For now, just generate stacks of diagonal matrices.
-    n = draw(integers(0, SQRT_MAX_ARRAY_SIZE),)
     stack_shape = draw(stack_shapes)
+    n = draw(integers(0, SQRT_MAX_ARRAY_SIZE // max(math.prod(stack_shape), 1)),)
     dtype = draw(dtypes)
     elements = one_of(
         from_dtype(dtype, min_value=0.5, allow_nan=False, allow_infinity=False),