Apply suggestions from code review

Co-authored-by: Lucas Colley <[email protected]>

Author: crusaderky

src/array_api_extra/testing.py

@@ -55,9 +55,9 @@ def lazy_xp_function(  # type: ignore[no-any-explicit]
     Tag a function to be tested on lazy backends.
 
     Tag a function, which must be imported in the test module globals, so that when any
-    tests defined in the same module are executed with `xp=jax.numpy` the function is
+    tests defined in the same module are executed with ``xp=jax.numpy`` the function is
     replaced with a jitted version of itself, and when it is executed with
-    `xp=dask.array` the function will raise if it attempts to materialize the graph.
+    ``xp=dask.array`` the function will raise if it attempts to materialize the graph.
     This will be later expanded to provide test coverage for other lazy backends.
 
     In order for the tag to be effective, the test or a fixture must call
@@ -69,7 +69,7 @@ def lazy_xp_function(  # type: ignore[no-any-explicit]
         Function to be tested.
     allow_dask_compute : int, optional
         Number of times `func` is allowed to internally materialize the Dask graph. This
-        is typically triggered by `bool()`, `float()`, or `np.asarray()`.
+        is typically triggered by ``bool()``, ``float()``, or ``np.asarray()``.
 
         Set to 1 if you are aware that `func` converts the input parameters to numpy and
         want to let it do so at least for the time being, knowing that it is going to be
@@ -78,16 +78,16 @@ def lazy_xp_function(  # type: ignore[no-any-explicit]
         If a test needs values higher than 1 to pass, it is a canary that the conversion
         to numpy/bool/float is happening multiple times, which translates to multiple
         computations of the whole graph. Short of making the function fully lazy, you
-        should at least add explicit calls to `np.asarray()` early in the function.
+        should at least add explicit calls to ``np.asarray()`` early in the function.
         *Note:* the counter of `allow_dask_compute` resets after each call to `func`, so
         a test function that invokes `func` multiple times should still work with this
         parameter set to 1.
 
         Default: 0, meaning that `func` must be fully lazy and never materialize the
         graph.
     jax_jit : bool, optional
-        Set to True to replace `func` with `jax.jit(func)` after calling the
-        :func:`patch_lazy_xp_functions` test helper with `xp=jax.numpy`. Set to False if
+        Set to True to replace `func` with ``jax.jit(func)`` after calling the
+        :func:`patch_lazy_xp_functions` test helper with ``xp=jax.numpy``. Set to False if
         `func` is only compatible with eager (non-jitted) JAX. Default: True.
     static_argnums : int | Sequence[int], optional
         Passed to jax.jit. Positional arguments to treat as static (compile-time
@@ -104,7 +104,7 @@ def lazy_xp_function(  # type: ignore[no-any-explicit]
 
     Examples
     --------
-    In `test_mymodule.py`::
+    In ``test_mymodule.py``::
 
       from array_api_extra.testing import lazy_xp_function from mymodule import myfunc
 
@@ -145,12 +145,12 @@ def patch_lazy_xp_functions(
     """
     Test lazy execution of functions tagged with :func:`lazy_xp_function`.
 
-    If `xp==jax.numpy`, search for all functions which have been tagged with
+    If ``xp==jax.numpy``, search for all functions which have been tagged with
     :func:`lazy_xp_function` in the globals of the module that defines the current test
     and wrap them with :func:`jax.jit`. Unwrap them at the end of the test.
 
-    If `xp==dask.array`, wrap the functions with a decorator that disables `compute()`
-    and `persist()`.
+    If ``xp==dask.array``, wrap the functions with a decorator that disables ``compute()``
+    and ``persist()``.
 
     This function should be typically called by your library's `xp` fixture that runs
     tests on multiple backends::
@@ -248,7 +248,7 @@ def _allow_dask_compute(
     msg = (
         f"Called `dask.compute()` or `dask.persist()` {n + 1} times, "
         f"but {n_str} calls are allowed. Set "
-        f"`lazy_xp_function({func_name}, allow_dask_compute={n + 1}) "
+        f"`lazy_xp_function({func_name}, allow_dask_compute={n + 1})` "
         "to allow for more (but note that this will harm performance). "
     )