From e89dba774e63d347413afb78aa8f629d82a9c5c0 Mon Sep 17 00:00:00 2001 From: nstarman Date: Tue, 27 Aug 2024 14:56:23 -0400 Subject: [PATCH 1/5] docs: update readme Signed-off-by: nstarman --- README.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/README.md b/README.md index 7e411a6..d427ed3 100644 --- a/README.md +++ b/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. - - [actions-badge]: https://github.com/GalacticDynamics/quaxed/workflows/CI/badge.svg [actions-link]: https://github.com/GalacticDynamics/quaxed/actions - [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 From a9e86aa149f02083e94a15248f470c1b7f6f9f2a Mon Sep 17 00:00:00 2001 From: nstarman Date: Tue, 27 Aug 2024 14:57:01 -0400 Subject: [PATCH 2/5] fix: vector_norm kwargs Signed-off-by: nstarman --- src/quaxed/array_api/linalg.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/src/quaxed/array_api/linalg.py b/src/quaxed/array_api/linalg.py index 7b5d5c8..bfa5f1c 100644 --- a/src/quaxed/array_api/linalg.py +++ b/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 @@ def vecdot(x1: ArrayLike, x2: ArrayLike, /, *, axis: int | None = None) -> Value @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) From cb4905eeb3cfd4bf92bf7ed7a9a36f02e3173441 Mon Sep 17 00:00:00 2001 From: nstarman Date: Tue, 27 Aug 2024 14:57:32 -0400 Subject: [PATCH 3/5] fix: _dot_general_p kwargs Signed-off-by: nstarman --- tests/myarray.py | 22 +++------------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/tests/myarray.py b/tests/myarray.py index 773e1f1..be2540c 100644 --- a/tests/myarray.py +++ b/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)) # ============================================================================== From 23b066162cd6d0271e37ea7b47d5905fcc1f1ccd Mon Sep 17 00:00:00 2001 From: nstarman Date: Tue, 27 Aug 2024 14:58:40 -0400 Subject: [PATCH 4/5] fix: floatX transfer Signed-off-by: nstarman --- src/quaxed/numpy/_core.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/quaxed/numpy/_core.py b/src/quaxed/numpy/_core.py index d135a08..b095aa9 100644 --- a/src/quaxed/numpy/_core.py +++ b/src/quaxed/numpy/_core.py @@ -462,6 +462,8 @@ def __dir__() -> list[str]: "euler_gamma", "flexible", "floating", + "float32", + "float64", "generic", "index_exp", "indices", From 5d2dc79473c0129988608ec35e947dc1a9e3c553 Mon Sep 17 00:00:00 2001 From: nstarman Date: Tue, 27 Aug 2024 15:15:51 -0400 Subject: [PATCH 5/5] build: docs Signed-off-by: nstarman --- pyproject.toml | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ddd135d..0cafe00 100644 --- a/pyproject.toml +++ b/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"