📝 docs: explain dynamic generation in docs (#119)

Signed-off-by: Nathaniel Starkman <[email protected]>

src/quaxed/__init__.py

@@ -1,6 +1,20 @@
-"""Copyright (c) 2023 Nathaniel Starkman. All rights reserved.
+"""Quaxified `jax.scipy`.
+
+This module wraps the functions in `jax.lax` with `quax.quaxify`. The wrapping
+happens dynamically through a module-level ``__dir__`` and ``__getattr__``. The
+list of available functions is in ``__all__`` and documented in the `jax.lax`
+library.
+
+In addition the following modules are supported:
+
+- `quaxed.lax.linalg`
+
+The contents of these modules are likewise dynamically wrapped with
+`quax.quaxify` and their contents is listed in their respective ``__all__`` and
+documented in their respective libraries.
+
+If a function is missing, please file an Issue.
 
-quaxed: Pre-quaxed libraries for multiple dispatch over abstract array types in JAX
 """
 # pylint: disable=C0415,W0621
 

src/quaxed/array_api/__init__.py

@@ -1,6 +1,8 @@
-"""Copyright (c) 2023 Nathaniel Starkman. All rights reserved.
+"""Quaxified `jax.experimental.array_api`.
+
+Note that this module is deprecated in newer versions of JAX and importing
+`quaxed.array_api` may therefore raise an `ImportError`.
 
-quaxed: Pre-quaxed libraries for multiple dispatch over abstract array types in JAX
 """
 
 # pylint: disable=redefined-builtin

src/quaxed/lax/__init__.py

@@ -1,4 +1,21 @@
-"""Quaxed :mod:`jax.lax`."""
+"""Quaxified `jax.lax`.
+
+This module wraps the functions in `jax.lax` with `quax.quaxify`. The wrapping
+happens dynamically through a module-level ``__dir__`` and ``__getattr__``. The
+list of available functions is in ``__all__`` and documented in the `jax.lax`
+library.
+
+In addition the following modules are supported:
+
+- `quaxed.lax.linalg`
+
+The contents of these modules are likewise dynamically wrapped with
+`quax.quaxify` and their contents is listed in their respective ``__all__`` and
+documented in their respective libraries.
+
+If a function is missing, please file an Issue.
+
+"""
 # pylint: disable=undefined-all-variable
 
 __all__ = [

src/quaxed/numpy/__init__.py

@@ -1,18 +1,38 @@
-"""Quaxed :mod:`jax.numpy`."""
+"""Quaxified `jax.numpy`.
+
+This module wraps the functions in `jax.numpy` with `quax.quaxify`. The wrapping
+happens dynamically through a module-level ``__dir__`` and ``__getattr__``. The
+list of available functions is in ``__all__`` and documented in the `jax.numpy`
+library.
+
+In addition the following modules are supported:
+
+- `quaxed.numpy.fft`
+- `quaxed.numpy.linalg`
+
+The contents of these modules are likewise dynamically wrapped with
+`quax.quaxify` and their contents is listed in their respective ``__all__`` and
+documented in their respective libraries.
+
+If a function is missing, please file an Issue.
+
+"""
 # pylint: disable=redefined-builtin
 
 from typing import Any
 
 from jaxtyping import install_import_hook
 
-from . import _core, _creation_functions, _higher_order
+from . import _core, _creation_functions, _higher_order, fft, linalg
 from ._creation_functions import *
 from ._higher_order import *
 
-__all__: list[str] = []
-__all__ += _core.__all__
-__all__ += _higher_order.__all__
-__all__ += _creation_functions.__all__
+__all__ = (
+    ["fft", "linalg"]  # noqa: RUF005
+    + _core.__all__
+    + _higher_order.__all__
+    + _creation_functions.__all__
+)
 
 
 # TODO: consolidate with ``_core.__getattr__``.

src/quaxed/operator.py

@@ -1,9 +1,9 @@
-"""Quaxed :external:`operator`.
+"""Quaxed `operator`.
 
-This module wraps the functions in :external:`operator` with
-:external:`quax.quaxify`. The wrapping happens dynamically through a
-module-level ``__dir__`` and ``__getattr__``. The list of available functions is
-in ``__all__`` and documented in the built-in :external:`operator` library.
+This module wraps the functions in `operator` with `quax.quaxify`. The wrapping
+happens dynamically through a module-level ``__dir__`` and ``__getattr__``. The
+list of available functions is in ``__all__`` and documented in the built-in
+`operator` library.
 
 """