♻️ refactor(quantity): ensure abstract quantities are abstract (#348)

src/unxt/_src/quantity/base.py

@@ -23,7 +23,7 @@
 
 from .api import is_unit_convertible, uconvert, ustrip
 from .mixins import AstropyQuantityCompatMixin, IPythonReprMixin, NumPyCompatMixin
-from unxt._src.units import unit as parse_unit, unit_of
+from unxt._src.units import unit_of
 from unxt._src.units.api import AbstractUnits
 
 if TYPE_CHECKING:
@@ -103,10 +103,10 @@ class AbstractQuantity(
 
     """
 
-    value: Shaped[Array, "*shape"] = eqx.field(converter=jax.numpy.asarray)
+    value: eqx.AbstractVar[Shaped[Array, "*shape"]]
     """The value of the `AbstractQuantity`."""
 
-    unit: AbstractUnits = eqx.field(static=True, converter=parse_unit)
+    unit: eqx.AbstractVar[AbstractUnits]
     """The unit associated with this value."""
 
     # ---------------------------------------------------------------

src/unxt/_src/quantity/base_parametric.py

@@ -8,8 +8,6 @@
 from typing import Any
 
 import equinox as eqx
-import jax
-import jax.core
 from astropy.units import PhysicalType, Unit
 from jaxtyping import Array, ArrayLike, Shaped
 from plum import dispatch, parametric, type_nonparametric, type_unparametrized
@@ -30,10 +28,10 @@ class AbstractParametricQuantity(AbstractQuantity):
 
     """
 
-    value: Shaped[Array, "*shape"] = eqx.field(converter=jax.numpy.asarray)
+    value: eqx.AbstractVar[Shaped[Array, "*shape"]]
     """The value of the `AbstractQuantity`."""
 
-    unit: Unit = eqx.field(static=True, converter=parse_unit)
+    unit: eqx.AbstractVar[Unit]
     """The unit associated with this value."""
 
     def __post_init__(self) -> None:

src/unxt/_src/quantity/core.py

@@ -6,11 +6,15 @@
 from dataclasses import replace
 from typing import final
 
-from jaxtyping import ArrayLike
+import equinox as eqx
+import jax
+from jaxtyping import Array, ArrayLike, Shaped
 from plum import parametric
 
 from .base import AbstractQuantity
 from .base_parametric import AbstractParametricQuantity
+from unxt._src.units import unit as parse_unit
+from unxt._src.units.api import AbstractUnits
 
 
 @final
@@ -102,6 +106,12 @@ class Quantity(AbstractParametricQuantity):
 
     """
 
+    value: Shaped[Array, "*shape"] = eqx.field(converter=jax.numpy.asarray)
+    """The value of the `AbstractQuantity`."""
+
+    unit: AbstractUnits = eqx.field(static=True, converter=parse_unit)
+    """The unit associated with this value."""
+
 
 @AbstractQuantity.__mod__.dispatch  # type: ignore[misc]
 def mod(self: Quantity["dimensionless"], other: ArrayLike) -> Quantity["dimensionless"]:

src/unxt/_src/quantity/fast.py

@@ -5,7 +5,13 @@
 
 from typing import Any
 
+import equinox as eqx
+import jax
+from jaxtyping import Array, Shaped
+
 from .base import AbstractQuantity
+from unxt._src.units import unit as parse_unit
+from unxt._src.units.api import AbstractUnits
 
 
 class UncheckedQuantity(AbstractQuantity):
@@ -14,6 +20,12 @@ class UncheckedQuantity(AbstractQuantity):
     This class is not parametrized by its dimensionality.
     """
 
+    value: Shaped[Array, "*shape"] = eqx.field(converter=jax.numpy.asarray)
+    """The value of the `AbstractQuantity`."""
+
+    unit: AbstractUnits = eqx.field(static=True, converter=parse_unit)
+    """The unit associated with this value."""
+
     def __class_getitem__(
         cls: type["UncheckedQuantity"], item: Any
     ) -> type["UncheckedQuantity"]:

src/unxt/_src/quantity/mixins.py

@@ -5,6 +5,7 @@
 from collections.abc import Callable, Sequence
 from typing import TYPE_CHECKING, Any, cast
 
+import equinox as eqx
 import numpy as np
 from astropy.units import CompositeUnit
 from jax.typing import ArrayLike
@@ -23,8 +24,8 @@
 class AstropyQuantityCompatMixin:
     """Mixin for compatibility with `astropy.units.Quantity`."""
 
-    value: ArrayLike
-    unit: AbstractUnits
+    value: eqx.AbstractVar[ArrayLike]
+    unit: eqx.AbstractVar[AbstractUnits]
     uconvert: Callable[[Any], "unxt.quantity.AbstractQuantity"]
     ustrip: Callable[[Any], ArrayLike]
 

tests/unit/test_quantity.py

@@ -14,7 +14,7 @@
 from hypothesis.extra.array_api import make_strategies_namespace
 from hypothesis.extra.numpy import array_shapes as np_array_shapes, arrays as np_arrays
 from jax.dtypes import canonicalize_dtype
-from jaxtyping import TypeCheckError
+from jaxtyping import Array, TypeCheckError
 from plum import convert, parametric
 
 import quaxed.numpy as jnp
@@ -504,6 +504,8 @@ def test_at():
 class NewQuantity(AbstractParametricQuantity):
     """Quantity with a flag."""
 
+    value: Array = eqx.field(converter=jnp.asarray)
+    unit: str = eqx.field(converter=u.unit)
     flag: bool = eqx.field(static=True, kw_only=True)