ENH: make MAX_ARRAY_SIZE configurable, lower the default to 1024

Changing the default is via the environment variable
ARRAY_API_TESTS_MAX_ARRAY_SIZE.

Author: ev-br

README.md

@@ -316,6 +316,14 @@ entries from the ``xfail.txt`` file instead of xfailing them. Anecdotally, we sa
 speed-ups by a factor of 4-5---which allowed us to use 4-5 larger values of
 ``--max-examples`` within the same time budget.
 
+#### Limiting the array sizes
+
+The test suite generates random arrays as inputs to functions it tests. "unvectorized"
+tests iterate over elements of arrays, which might be slow. If the run time becomes
+a problem, you can limit the maximum number of elements in generated arrays by
+setting the environment variable ``ARRAY_API_TESTS_MAX_ARRAY_SIZE`` to the
+desired value. By default, it is set to 1024.
+
 
 ## Contributing
 

array_api_tests/hypothesis_helpers.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import re
 from contextlib import contextmanager
 from functools import wraps
@@ -232,7 +233,7 @@ def all_floating_dtypes() -> SearchStrategy[DataType]:
     lambda i: getattr(xp, i))
 
 # Limit the total size of an array shape
-MAX_ARRAY_SIZE = 10000
+MAX_ARRAY_SIZE = int(os.environ.get("ARRAY_API_TESTS_MAX_ARRAY_SIZE", 1024))
 # Size to use for 2-dim arrays
 SQRT_MAX_ARRAY_SIZE = int(math.sqrt(MAX_ARRAY_SIZE))
 

array_api_tests/test_creation_functions.py

@@ -499,7 +499,7 @@ def test_meshgrid(dtype, data):
     shapes = data.draw(
         st.integers(1, 5).flatmap(
             lambda n: hh.mutually_broadcastable_shapes(
-                n, min_dims=1, max_dims=1, max_side=5
+                n, min_dims=1, max_dims=1, max_side=4
             )
         ),
         label="shapes",
@@ -509,7 +509,7 @@ def test_meshgrid(dtype, data):
         x = data.draw(hh.arrays(dtype=dtype, shape=shape), label=f"x{i}")
         arrays.append(x)
     # sanity check
-    assert math.prod(math.prod(x.shape) for x in arrays) <= hh.MAX_ARRAY_SIZE
+    # assert math.prod(math.prod(x.shape) for x in arrays) <= hh.MAX_ARRAY_SIZE
     out = xp.meshgrid(*arrays)
     for i, x in enumerate(out):
         ph.assert_dtype("meshgrid", in_dtype=dtype, out_dtype=x.dtype, repr_name=f"out[{i}].dtype")