fix: upgrade type hints.

By dropping support for Python 3.9, it is possible to replace the
`Union` and `Optional` type hints with the `|` operator. At the
same time, it is also possible to simplify some of the jaxtyped
type hints.

BREAKING CHANGES: Sequence support for `_atleast_2d_consistent`
has been dropped. The functionality was unused and significantly
complicated the type hints.

Author: tm91236

CHANGELOG.md

@@ -17,11 +17,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
- -
+ - Type hints now use python 3.10 syntax; `Union` and `Optional` types have been replaced with `A | B` and `A | None`.
+ - Jaxtyped arguments now have simpler/stricter shape requirements;
 
 ### Removed
 
  - Internal `SilentTQDM` class has been removed; all user-facing functionality has been preserved.
+ - Internal `_atleast_2d_consistent` function no longer supports sequences of arrays; all user-facing functionality has been preserved.
  - Dropped support for python 3.9 which is near end-of-life.
 
 ### Deprecated

benchmark/blobs_benchmark.py

@@ -31,13 +31,12 @@
 import json
 import os
 import time
-from typing import Union
 
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
 import numpy as np
-from jaxtyping import Array, Shaped
+from jaxtyping import Array, ArrayLike, Scalar, Shaped
 from sklearn.datasets import make_blobs
 
 from coreax import Data
@@ -95,25 +94,21 @@ def setup_stein_kernel(
     kde = jsp.stats.gaussian_kde(dataset.data[idx].T)
 
     # Define the score function as the gradient of log density given by the KDE
-    def score_function(
-        x: Union[Shaped[Array, " n d"], Shaped[Array, ""], float, int],
-    ) -> Union[Shaped[Array, " n d"], Shaped[Array, " 1 1"]]:
+    def score_function(x: Shaped[ArrayLike, " n d"]) -> Shaped[Array, " n d"]:
         """
         Compute the score function (gradient of log density) for a single point.
 
         :param x: Input point represented as array
         :return: Gradient of log probability density at the given point
         """
 
-        def logpdf_single(x: Shaped[Array, " d"]) -> Shaped[Array, ""]:
+        def logpdf_single(x: Shaped[Array, " d"]) -> Scalar:
             return kde.logpdf(x.reshape(1, -1))[0]
 
-        return jax.grad(logpdf_single)(x)
+        _x = jnp.asarray(x)
+        return jax.grad(logpdf_single)(_x)
 
-    return SteinKernel(
-        base_kernel=sq_exp_kernel,
-        score_function=score_function,
-    )
+    return SteinKernel(base_kernel=sq_exp_kernel, score_function=score_function)
 
 
 def setup_solvers(

benchmark/david_benchmark.py

@@ -38,7 +38,6 @@
 import math
 import time
 from pathlib import Path
-from typing import Optional
 
 import equinox as eqx
 import jax.numpy as jnp
@@ -60,7 +59,7 @@
 # (disable line too long and too many statements ruff)
 def benchmark_coreset_algorithms(
     in_path: Path = Path("../examples/data/david_orig.png"),
-    out_path: Optional[Path] = Path(
+    out_path: Path | None = Path(
         "../examples/benchmarking_images/david_benchmark_results.png"
     ),
     downsampling_factor: int = 1,

benchmark/mnist_benchmark.py

@@ -40,7 +40,7 @@
 import json
 import os
 import time
-from typing import Any, NamedTuple, Optional, Union
+from typing import Any, NamedTuple
 
 import equinox as eqx
 import jax
@@ -143,7 +143,7 @@ def __call__(self, x: jnp.ndarray, training: bool = True) -> jnp.ndarray:
 class TrainState(train_state.TrainState):
     """Custom train state with batch statistics and dropout RNG."""
 
-    batch_stats: Optional[dict[str, jnp.ndarray]]
+    batch_stats: dict[str, jnp.ndarray] | None
     dropout_rng: KeyArrayLike
 
 
@@ -416,7 +416,7 @@ def prepare_datasets() -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarr
 def train_model(
     data_bundle: dict[str, jnp.ndarray],
     key: KeyArrayLike,
-    config: dict[str, Union[int, float]],
+    config: dict[str, int | float],
 ) -> dict[str, float]:
     """
     Train the model and return the results.

coreax/approximation.py

@@ -30,14 +30,13 @@
 
 from collections.abc import Callable
 from functools import partial
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Literal
 
 import jax
 import jax.numpy as jnp
 import jax.random as jr
-from jax import Array
-from jaxtyping import Shaped
-from typing_extensions import Literal, override
+from jaxtyping import Array, ArrayLike, Shaped
+from typing_extensions import override
 
 from coreax.data import Data, _atleast_2d_consistent
 from coreax.kernels import UniCompositeKernel
@@ -176,16 +175,7 @@ class MonteCarloApproximateKernel(RandomRegressionKernel):
     """
 
     def gramian_row_mean(
-        self,
-        x: Union[
-            Shaped[Array, " n d"],
-            Shaped[Array, " d"],
-            Shaped[Array, ""],
-            float,
-            int,
-            Data,
-        ],
-        **kwargs: Any,
+        self, x: Shaped[ArrayLike | Data, " n d"], **kwargs: Any
     ) -> Shaped[Array, " n"]:
         r"""
         Approximate the Gramian row-mean by Monte-Carlo sampling.
@@ -231,16 +221,7 @@ class ANNchorApproximateKernel(RandomRegressionKernel):
     """
 
     def gramian_row_mean(
-        self,
-        x: Union[
-            Shaped[Array, " n d"],
-            Shaped[Array, " d"],
-            Shaped[Array, ""],
-            float,
-            int,
-            Data,
-        ],
-        **kwargs: Any,
+        self, x: Shaped[ArrayLike | Data, " n d"], **kwargs: Any
     ) -> Shaped[Array, " n"]:
         r"""
         Approximate the Gramian row-mean by random regression on ANNchor points.
@@ -304,16 +285,7 @@ class NystromApproximateKernel(RandomRegressionKernel):
     """
 
     def gramian_row_mean(
-        self,
-        x: Union[
-            Shaped[Array, " n d"],
-            Shaped[Array, " d"],
-            Shaped[Array, ""],
-            float,
-            int,
-            Data,
-        ],
-        **kwargs: Any,
+        self, x: Shaped[ArrayLike | Data, " n d"], **kwargs: Any
     ) -> Shaped[Array, " n"]:
         r"""
         Approximate the Gramian row-mean by Nystrom approximation.

coreax/benchmark_util.py

@@ -21,13 +21,13 @@
 """
 
 from collections.abc import Callable
-from typing import Optional, TypeVar, Union
+from typing import TypeVar
 
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
 import numpy as np
-from jaxtyping import Array, Float, Shaped
+from jaxtyping import Array, ArrayLike, Float, Scalar, Shaped
 
 from coreax import Coresubset, Data, SupervisedData
 from coreax.kernels import SquaredExponentialKernel, SteinKernel, median_heuristic
@@ -61,12 +61,12 @@ class IterativeKernelHerding(KernelHerding[_Data]):  # pylint: disable=too-many-
     """
 
     num_iterations: int = 1
-    t_schedule: Optional[Array] = None
+    t_schedule: Array | None = None
 
     def reduce(
         self,
         dataset: _Data,
-        solver_state: Optional[HerdingState] = None,
+        solver_state: HerdingState | None = None,
     ) -> tuple[Coresubset[_Data], HerdingState]:
         """
         Perform Kernel Herding reduction followed by additional refinement iterations.
@@ -118,7 +118,7 @@ def initialise_solvers(  # noqa: C901
     train_data_umap: Data,
     key: KeyArrayLike,
     cpp_oversampling_factor: int,
-    leaf_size: Optional[int] = None,
+    leaf_size: int | None = None,
 ) -> dict[str, Callable[[int], Solver]]:
     """
     Initialise and return a list of solvers for various coreset algorithms.
@@ -147,7 +147,7 @@ def initialise_solvers(  # noqa: C901
     kernel = SquaredExponentialKernel(length_scale=length_scale)
     sqrt_kernel = kernel.get_sqrt_kernel(16)
 
-    def _get_thinning_solver(_size: int) -> Union[KernelThinning, MapReduce]:
+    def _get_thinning_solver(_size: int) -> KernelThinning | MapReduce:
         """
         Set up kernel thinning solver.
 
@@ -169,7 +169,7 @@ def _get_thinning_solver(_size: int) -> Union[KernelThinning, MapReduce]:
             return thinning_solver
         return MapReduce(thinning_solver, leaf_size=leaf_size)
 
-    def _get_herding_solver(_size: int) -> Union[KernelHerding, MapReduce]:
+    def _get_herding_solver(_size: int) -> KernelHerding | MapReduce:
         """
         Set up kernel herding solver.
 
@@ -185,7 +185,7 @@ def _get_herding_solver(_size: int) -> Union[KernelHerding, MapReduce]:
             return herding_solver
         return MapReduce(herding_solver, leaf_size=leaf_size)
 
-    def _get_stein_solver(_size: int) -> Union[SteinThinning, MapReduce]:
+    def _get_stein_solver(_size: int) -> SteinThinning | MapReduce:
         """
         Set up Stein thinning solver.
 
@@ -199,20 +199,19 @@ def _get_stein_solver(_size: int) -> Union[SteinThinning, MapReduce]:
         kde = jsp.stats.gaussian_kde(train_data_umap.data[idx].T)
 
         # Define the score function as the gradient of log density given by the KDE
-        def score_function(
-            x: Union[Shaped[Array, " n d"], Shaped[Array, ""], float, int],
-        ) -> Union[Shaped[Array, " n d"], Shaped[Array, " 1 1"]]:
+        def score_function(x: Shaped[ArrayLike, " n d"]) -> Shaped[Array, " n d"]:
             """
             Compute the score function (gradient of log density) for a single point.
 
             :param x: Input point represented as array.
             :return: Gradient of log probability density at the given point.
             """
 
-            def logpdf_single(x: Shaped[Array, " d"]) -> Shaped[Array, ""]:
+            def logpdf_single(x: Shaped[Array, " d"]) -> Scalar:
                 return kde.logpdf(x.reshape(1, -1))[0]
 
-            return jax.grad(logpdf_single)(x)
+            _x = jnp.asarray(x)
+            return jax.grad(logpdf_single)(_x)
 
         stein_kernel = SteinKernel(
             base_kernel=kernel,
@@ -264,7 +263,7 @@ def _get_compress_solver(_size: int) -> CompressPlusPlus:
 
     def _get_probabilistic_herding_solver(
         _size: int,
-    ) -> Union[IterativeKernelHerding, MapReduce]:
+    ) -> IterativeKernelHerding | MapReduce:
         """
         Set up KernelHerding with probabilistic selection.
 
@@ -289,7 +288,7 @@ def _get_probabilistic_herding_solver(
 
     def _get_iterative_herding_solver(
         _size: int,
-    ) -> Union[IterativeKernelHerding, MapReduce]:
+    ) -> IterativeKernelHerding | MapReduce:
         """
         Set up KernelHerding with probabilistic selection.
 
@@ -314,7 +313,7 @@ def _get_iterative_herding_solver(
 
     def _get_cubic_iterative_herding_solver(
         _size: int,
-    ) -> Union[IterativeKernelHerding, MapReduce]:
+    ) -> IterativeKernelHerding | MapReduce:
         """
         Set up KernelHerding with probabilistic selection.
 

coreax/coreset.py

@@ -20,13 +20,12 @@
     Final,
     Generic,
     TypeVar,
-    Union,
     overload,
 )
 
 import equinox as eqx
 import jax.numpy as jnp
-from jaxtyping import Array, Shaped
+from jaxtyping import Array, Scalar, Shaped
 from typing_extensions import Self, override
 
 from coreax.data import Data, SupervisedData, as_data
@@ -71,9 +70,7 @@ def pre_coreset_data(self) -> _TOriginalData_co:
     def solve_weights(self, solver: WeightsOptimiser[Data], **solver_kwargs) -> Self:
         """Return a copy of 'self' with weights solved by 'solver'."""
 
-    def compute_metric(
-        self, metric: Metric[Data], **metric_kwargs
-    ) -> Shaped[Array, ""]:
+    def compute_metric(self, metric: Metric[Data], **metric_kwargs) -> Scalar:
         """Return metric-distance between `self.pre_coreset_data` and `self.coreset`."""
         return metric.compute(self.pre_coreset_data, self.points, **metric_kwargs)
 
@@ -131,30 +128,30 @@ def __init__(self, nodes: Data, pre_coreset_data: _TOriginalData_co) -> None:
     @classmethod
     @overload
     def build(
-        cls, nodes: Union[Data, Array], pre_coreset_data: Array
+        cls, nodes: Data | Array, pre_coreset_data: Array
     ) -> "PseudoCoreset[Data]": ...
 
     @classmethod
     @overload
     def build(
         cls,
-        nodes: Union[Data, Array],
+        nodes: Data | Array,
         pre_coreset_data: tuple[Array, Array],
     ) -> "PseudoCoreset[SupervisedData]": ...
 
     @classmethod
     @overload
     def build(
         cls,
-        nodes: Union[Data, Array],
+        nodes: Data | Array,
         pre_coreset_data: _TOriginalData,
     ) -> "PseudoCoreset[_TOriginalData]": ...
 
     @classmethod
     def build(
         cls,
-        nodes: Union[Data, Array],
-        pre_coreset_data: Union[_TOriginalData, Array, tuple[Array, Array]],
+        nodes: Data | Array,
+        pre_coreset_data: _TOriginalData | Array | tuple[Array, Array],
     ) -> "PseudoCoreset[Data]\
         | PseudoCoreset[SupervisedData]\
         | PseudoCoreset[_TOriginalData]\
@@ -253,30 +250,30 @@ def __init__(self, indices: Data, pre_coreset_data: _TOriginalData_co) -> None:
     @classmethod
     @overload
     def build(
-        cls, indices: Union[Data, Array], pre_coreset_data: Array
+        cls, indices: Data | Array, pre_coreset_data: Array
     ) -> "Coresubset[Data]": ...
 
     @classmethod
     @overload
     def build(
         cls,
-        indices: Union[Data, Array],
+        indices: Data | Array,
         pre_coreset_data: tuple[Array, Array],
     ) -> "Coresubset[SupervisedData]": ...
 
     @classmethod
     @overload
     def build(
         cls,
-        indices: Union[Data, Array],
+        indices: Data | Array,
         pre_coreset_data: _TOriginalData,
     ) -> "Coresubset[_TOriginalData]": ...
 
     @classmethod
     def build(
         cls,
-        indices: Union[Data, Array],
-        pre_coreset_data: Union[_TOriginalData, Array, tuple[Array, Array]],
+        indices: Data | Array,
+        pre_coreset_data: _TOriginalData | Array | tuple[Array, Array],
     ) -> "Coresubset[Data] | Coresubset[SupervisedData] | Coresubset[_TOriginalData]":
         """
         Construct a Coresubset from Data or raw Arrays.