feat: dep quaxed

pyproject.toml

@@ -23,7 +23,7 @@
     "Typing :: Typed",
   ]
   dependencies = [
-    "array_api_jax_compat >= 0.1",
+    "quaxed >= 0.2",
     "astropy",
     "equinox",
     "jax",
@@ -134,7 +134,7 @@
     ignore_missing_imports = true
     module = [
       "array_api.*",
-      "array_api_jax_compat.*",
+      "quaxed.*",
       "astropy.*",
       "equinox.*",
       "hypothesis.*",
@@ -188,7 +188,7 @@
 
     [tool.ruff.lint.isort]
       combine-as-imports = true
-      known-first-party = ["array_api_jax_compat", "jax_quantity"]
+      known-first-party = ["quaxed", "jax_quantity"]
       known-local-folder = ["coordinax"]
 
 

src/coordinax/_base.py

@@ -20,7 +20,7 @@
 from jaxtyping import ArrayLike
 from plum import dispatch
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from ._utils import classproperty, dataclass_items, dataclass_values, full_shaped

src/coordinax/_checks.py

@@ -5,7 +5,7 @@
 
 import equinox as eqx
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from coordinax._typing import BatchableAngle, BatchableLength

src/coordinax/_d1/builtin.py

@@ -16,7 +16,7 @@
 import equinox as eqx
 import jax
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import Abstract1DVector, Abstract1DVectorDifferential

src/coordinax/_d1/compat.py

@@ -6,7 +6,7 @@
 from jaxtyping import Shaped
 from plum import conversion_method
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import Abstract1DVector

src/coordinax/_d2/builtin.py

@@ -16,7 +16,7 @@
 import equinox as eqx
 import jax
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import Abstract2DVector, Abstract2DVectorDifferential

src/coordinax/_d2/compat.py

@@ -6,7 +6,7 @@
 from jaxtyping import Shaped
 from plum import conversion_method
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import Abstract2DVector

src/coordinax/_d2/transform.py

@@ -6,7 +6,7 @@
 
 from plum import dispatch
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 
 from .base import Abstract2DVector, Abstract2DVectorDifferential
 from .builtin import (

src/coordinax/_d3/__init__.py

@@ -1,5 +1,5 @@
 # pylint: disable=duplicate-code
-"""3-dimensional."""
+"""3-dimensional representations."""
 
 from . import base, builtin, compat, operate, transform
 from .base import *

src/coordinax/_d3/builtin.py

@@ -18,7 +18,7 @@
 import equinox as eqx
 import jax
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import Abstract3DVector, Abstract3DVectorDifferential

src/coordinax/_d3/compat.py

@@ -8,7 +8,7 @@
 from jaxtyping import Shaped
 from plum import conversion_method, convert
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import Abstract3DVector

src/coordinax/_d3/transform.py

@@ -6,7 +6,7 @@
 
 from plum import dispatch
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 
 from .base import Abstract3DVector, Abstract3DVectorDifferential
 from .builtin import (

src/coordinax/_d4/compat.py

@@ -6,7 +6,7 @@
 from jaxtyping import Shaped
 from plum import conversion_method, convert
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .spacetime import FourVector

src/coordinax/_d4/spacetime.py

@@ -12,7 +12,7 @@
 import jax.numpy as jnp
 from jaxtyping import Shaped
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import Abstract4DVector

src/coordinax/_transform.py

@@ -1,4 +1,4 @@
-"""Representation of coordinates in different systems."""
+"""Transformations between representations."""
 
 __all__ = ["represent_as"]
 
@@ -10,7 +10,7 @@
 import jax
 from plum import dispatch
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from ._base import AbstractVector, AbstractVectorDifferential

src/coordinax/_utils.py

@@ -14,7 +14,7 @@
     runtime_checkable,
 )
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 
 if TYPE_CHECKING:
     from coordinax._base import AbstractVectorBase

src/coordinax/operators/_galilean/boost.py

@@ -11,7 +11,7 @@
 import jax.numpy as jnp
 from plum import convert
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import AbstractGalileanOperator
@@ -47,7 +47,7 @@ class GalileanBoostOperator(AbstractGalileanOperator):
     --------
     We start with the required imports:
 
-    >>> import array_api_jax_compat as xp
+    >>> import quaxed.array_api as xp
     >>> from jax_quantity import Quantity
     >>> from coordinax import CartesianDifferential3D, Cartesian3DVector
     >>> import coordinax.operators as co

src/coordinax/operators/_galilean/composite.py

@@ -8,7 +8,7 @@
 
 import equinox as eqx
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import AbstractGalileanOperator

src/coordinax/operators/_galilean/rotation.py

@@ -14,7 +14,7 @@
 from plum import convert
 from quax import quaxify
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import AbstractGalileanOperator
@@ -145,7 +145,7 @@ def is_inertial(self) -> Literal[True]:
 
         Examples
         --------
-        >>> import array_api_jax_compat as xp
+        >>> import quaxed.array_api as xp
         >>> from jax_quantity import Quantity
         >>> from coordinax.operators import GalileanRotationOperator
 
@@ -166,7 +166,7 @@ def inverse(self) -> "GalileanRotationOperator":
 
         Examples
         --------
-        >>> import array_api_jax_compat as xp
+        >>> import quaxed.array_api as xp
         >>> from jax_quantity import Quantity
         >>> from coordinax.operators import GalileanRotationOperator
 
@@ -194,7 +194,7 @@ def __call__(
 
         Examples
         --------
-        >>> import array_api_jax_compat as xp
+        >>> import quaxed.array_api as xp
         >>> from jax_quantity import Quantity
         >>> from coordinax import Cartesian3DVector, CartesianDifferential3D
         >>> from coordinax.operators import GalileanRotationOperator
@@ -226,7 +226,7 @@ def __call__(
 
         Examples
         --------
-        >>> import array_api_jax_compat as xp
+        >>> import quaxed.array_api as xp
         >>> from jax_quantity import Quantity
         >>> from coordinax import Cartesian3DVector, CartesianDifferential3D
         >>> from coordinax.operators import GalileanRotationOperator

src/coordinax/operators/_galilean/translation.py

@@ -11,7 +11,7 @@
 import jax.numpy as jnp
 from plum import convert
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .base import AbstractGalileanOperator
@@ -79,7 +79,7 @@ class GalileanSpatialTranslationOperator(AbstractGalileanOperator):
     --------
     We start with the required imports:
 
-    >>> import array_api_jax_compat as xp
+    >>> import quaxed.array_api as xp
     >>> from jax_quantity import Quantity
     >>> import coordinax as cx
 
@@ -324,7 +324,7 @@ class GalileanTranslationOperator(AbstractGalileanOperator):
     --------
     We start with the required imports:
 
-    >>> import array_api_jax_compat as xp
+    >>> import quaxed.array_api as xp
     >>> from jax_quantity import Quantity
     >>> import coordinax.operators as co
 

tests/test_base.py

@@ -10,7 +10,7 @@
 import pytest
 from quax import quaxify
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from coordinax import (

tests/test_d2.py

@@ -5,7 +5,7 @@
 import pytest
 from quax import quaxify
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .test_base import AbstractVectorDifferentialTest, AbstractVectorTest, array_equal

tests/test_d3.py

@@ -8,7 +8,7 @@
 from astropy.coordinates.tests.test_representation import representation_equal
 from plum import convert
 
-import array_api_jax_compat as xp
+import quaxed.array_api as xp
 from jax_quantity import Quantity
 
 from .test_base import AbstractVectorDifferentialTest, AbstractVectorTest, array_equal