Subclasses for Arrayish Objects (#120)

* ➕ dep-add(optype): for precise type definitions
* ✨ feat(arrayish): mixins for dunder methods, to support lax and numpy operations on arrayish objects
* 📝 docs(arrayish): add module to docs

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

Author: nstarman

.pre-commit-config.yaml

@@ -81,7 +81,7 @@ repos:
       - id: mypy
         files: src
         additional_dependencies:
-          - pytest
+          - optype
         exclude: |
           (?x)^(
             src/quaxed/lax/__init__.py|
@@ -111,4 +111,4 @@ repos:
         name: Disallow improper capitalization
         language: pygrep
         entry: PyBind|Numpy|Cmake|CCache|Github|PyTest
-        exclude: .pre-commit-config.yaml
+        exclude: .pre-commit-config.yaml|src/quaxed/experimental/arrayish.py|src/quaxed/experimental/_arrayish

docs/api/arrayish.md

@@ -0,0 +1,3 @@
+# quaxed.experimental.arrayish
+
+::: quaxed.experimental.arrayish

mkdocs.yml

@@ -102,3 +102,4 @@ nav:
       - "api/scipy.md"
       - "api/array_api.md"
       - "api/lax.md"
+      - "api/arrayish.md"

pyproject.toml

@@ -26,9 +26,10 @@
     "jax>=0.4.3",
     "jaxlib>=0.4.3",
     "jaxtyping>=0.2.34",
+    "optype>=0.8.0",
     "plum-dispatch>=2.5.2",
     "quax>=0.0.5",
-  ]
+]
 
   [project.urls]
     Homepage = "https://github.com/GalacticDynamics/quaxed"

src/quaxed/__init__.py

@@ -1,45 +1,40 @@
-"""Quaxified `jax.scipy`.
+"""Pre-`quaxify`ed jax and related libraries.
 
-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.
+`quax` is JAX + multiple dispatch + custom array-ish objects. `quaxed` is a
+drop-in replacement for many JAX and related libraries that applies
+`quax.quaxify` to the original JAX functions, enabling custom array-ish objects
+to be used with those functions, not only jax arrays.
 
 """
-# pylint: disable=C0415,W0621
-
-from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
-from . import _jax, lax, numpy, scipy
-from ._jax import *
+__all__ = [
+    # Modules
+    "lax",
+    "numpy",
+    "scipy",
+    "experimental",
+    # Jax functions
+    "device_put",
+    "grad",
+    "hessian",
+    "jacfwd",
+    "jacrev",
+    "value_and_grad",
+]
+
+from typing import TYPE_CHECKING, Any
+
+from . import experimental, lax, numpy, scipy
+from ._jax import device_put, grad, hessian, jacfwd, jacrev, value_and_grad
 from ._setup import JAX_VERSION
 from ._version import version as __version__  # noqa: F401
 
-__all__ = ["lax", "numpy", "scipy"]
-__all__ += _jax.__all__
-
 if JAX_VERSION < (0, 4, 32):
     from . import array_api
 
     __all__ += ["array_api"]
 
 
-if TYPE_CHECKING:
-    from typing import Any
-
-
 def __getattr__(name: str) -> Any:  # TODO: fuller annotation
     """Forward all other attribute accesses to Quaxified JAX."""
     import sys
@@ -59,4 +54,4 @@ def __getattr__(name: str) -> Any:  # TODO: fuller annotation
 
 
 # Clean up the namespace
-del TYPE_CHECKING
+del TYPE_CHECKING, Any

src/quaxed/experimental/__init__.py

@@ -0,0 +1 @@
+"""Experimental modules."""

src/quaxed/experimental/_arrayish/__init__.py

@@ -0,0 +1,3 @@
+"""Arrayish."""
+
+__all__: list[str] = []