fix: misc

README.md

@@ -53,13 +53,9 @@ If you found this library to be useful and want to support the development and
 maintenance of lower-level utility libraries for the scientific community,
 please consider citing this work.
 
-<!-- SPHINX-START -->
-
 <!-- prettier-ignore-start -->
 [actions-badge]:            https://github.com/GalacticDynamics/quaxed/workflows/CI/badge.svg
 [actions-link]:             https://github.com/GalacticDynamics/quaxed/actions
-<!-- [github-discussions-badge]: https://img.shields.io/static/v1?label=Discussions&message=Ask&color=blue&logo=github
-[github-discussions-link]:  https://github.com/GalacticDynamics/quaxed/discussions -->
 [pypi-link]:                https://pypi.org/project/quaxed/
 [pypi-platforms]:           https://img.shields.io/pypi/pyversions/quaxed
 [pypi-version]:             https://img.shields.io/pypi/v/quaxed

pyproject.toml

@@ -42,21 +42,19 @@ test = [
     "pytest-github-actions-annotate-failures",  # only applies to GH Actions
 ]
 docs = [
-    "mkdocs==1.3.0",  # Main documentation generator.
-    "mkdocs-material==7.3.6",  # Theme
+    "griffe < 1.0",  # For Python structure signatures"
+    "mkdocs==1.6.0",  # Main documentation generator.
+    "mkdocs-material==9.5",  # Theme
     "mkdocs_include_exclude_files==0.0.1",  # Tweak which files are included/excluded
     "mkdocstrings[python] >= 0.18",  # Autogenerate documentation from docstrings.
-    "mknotebooks==0.7.1",  # Turn Jupyter Lab notebooks into webpages.
-    "nbconvert==6.5.0",
-    "pygments==2.14.0",
-    "pymdown-extensions==9.4",  # Markdown extensions e.g. to handle LaTeX.
+    "mknotebooks==0.8",  # Turn Jupyter Lab notebooks into webpages.
+    "nbconvert==7.16",
+    "pygments==2.16",
+    "pymdown-extensions==10.2",  # Markdown extensions e.g. to handle LaTeX.
     "pytkdocs_tweaks==0.0.8",  # Tweaks mkdocstrings to improve various aspects
-    "jinja2==3.0.3"            # Older version. After 3.1.0 seems to be incompatible with current versions of mkdocstrings.
-]
-dev = [
-    "quaxed[test]",
-    "quaxed[docs]",
+    "jinja2==3.1"
 ]
+dev = ["quaxed[test,docs]"]
 
 [project.urls]
 Homepage = "https://github.com/GalacticDynamics/quaxed"

src/quaxed/array_api/linalg.py

@@ -28,7 +28,7 @@
 
 
 from collections.abc import Sequence
-from typing import Literal
+from typing import Any, Literal
 
 import jax.numpy as jnp
 from jax.experimental import array_api
@@ -167,12 +167,5 @@
 
 
 @quaxify
-def vector_norm(
-    x: ArrayLike,
-    /,
-    *,
-    axis: int | tuple[int, ...] | None = None,
-    keepdims: bool = False,
-    ord: int | float = 2,  # pylint: disable=redefined-builtin
-) -> Value:
-    return array_api.linalg.vector_norm(x, axis=axis, keepdims=keepdims, ord=ord)
+def vector_norm(x: ArrayLike, /, **kwargs: Any) -> Value:
+    return array_api.linalg.vector_norm(x, **kwargs)

src/quaxed/numpy/_core.py

@@ -462,6 +462,8 @@ def __dir__() -> list[str]:
         "euler_gamma",
         "flexible",
         "floating",
+        "float32",
+        "float64",
         "generic",
         "index_exp",
         "indices",

tests/myarray.py

@@ -8,9 +8,8 @@
 import jax
 import jax.experimental.array_api as jax_xp
 from jax import Device, lax
-from jax._src.lax.lax import DotDimensionNumbers, PrecisionLike
 from jax._src.lax.slicing import GatherDimensionNumbers, GatherScatterMode
-from jax._src.typing import DTypeLike, Shape
+from jax._src.typing import Shape
 from jaxtyping import ArrayLike
 from quax import ArrayValue, register
 
@@ -424,23 +423,8 @@ def _div_p(x: MyArray, y: ArrayLike) -> MyArray:
 
 
 @register(lax.dot_general_p)  # TODO: implement
-def _dot_general_p(
-    lhs: MyArray,
-    rhs: MyArray,
-    *,
-    dimension_numbers: DotDimensionNumbers,
-    precision: PrecisionLike = None,
-    preferred_element_type: DTypeLike | None = None,
-) -> MyArray:
-    return MyArray(
-        lax.dot_general_p.bind(
-            lhs.array,
-            rhs.array,
-            dimension_numbers=dimension_numbers,
-            precision=precision,
-            preferred_element_type=preferred_element_type,
-        ),
-    )
+def _dot_general_p(lhs: MyArray, rhs: MyArray, **kwargs: Any) -> MyArray:
+    return MyArray(lax.dot_general_p.bind(lhs.array, rhs.array, **kwargs))
 
 
 # ==============================================================================