Merge pull request #627 from lucascolley/lazy-test

TST/DOC/DEV: add missing `lazy_xp_function` calls, update contributing docs

Author: lucascolley

docs/contributing.md

@@ -20,6 +20,8 @@ Thanks to [all contributors](contributors.md) so far!
 - Import your function to `src/array_api_extra/__init__.py` and add it to
   `__all__` there.
 - Add a test class for your function in `tests/test_funcs.py`.
+  - Ensure that `lazy_xp_function` is called on the function if lazy backends
+    are supposed to be tested.
 - Add your function to `docs/api-reference.md`.
 - [Make a PR!](https://github.com/data-apis/array-api-extra/pulls)
 

src/array_api_extra/_delegation.py

@@ -1128,6 +1128,15 @@ def union1d(a: Array, b: Array, /, *, xp: ModuleType | None = None) -> Array:
     -------
     Array
         Unique, sorted union of the input arrays.
+
+    See Also
+    --------
+    jax.numpy.union1d : Corresponding function in JAX.
+
+    Notes
+    -----
+    This function is not compatible with `jax.jit`.
+    See the docstring of the corresponding JAX function for more information.
     """
     if xp is None:
         xp = array_namespace(a, b)

tests/test_funcs.py

@@ -51,20 +51,27 @@
 from array_api_extra.testing import lazy_xp_function
 
 lazy_xp_function(apply_where)
+lazy_xp_function(argpartition)
 lazy_xp_function(atleast_nd)
+lazy_xp_function(broadcast_shapes)
 lazy_xp_function(cov)
 lazy_xp_function(create_diagonal)
+lazy_xp_function(default_dtype)
 lazy_xp_function(expand_dims)
+lazy_xp_function(isclose)
+lazy_xp_function(isin)
 lazy_xp_function(kron)
 lazy_xp_function(nan_to_num)
 lazy_xp_function(nunique)
 lazy_xp_function(one_hot)
 lazy_xp_function(pad)
+lazy_xp_function(partition)
 # FIXME calls in1d which calls xp.unique_values without size
 lazy_xp_function(setdiff1d, jax_jit=False)
+lazy_xp_function(sinc)
+lazy_xp_function(union1d, jax_jit=False)
 lazy_xp_function(xpx_searchsorted)
 lazy_xp_function(_funcs_searchsorted)
-lazy_xp_function(sinc)
 
 
 class TestApplyWhere: