Add GPU-accelerated CPA implementation

project-words.txt

@@ -7,6 +7,7 @@ Babuska
 Benadjila
 Bursztein
 CHES
+CUDA
 Cassiers
 Chmielewski
 DSOWINDOW
@@ -68,6 +69,7 @@ allclose
 arange
 argmax
 argmin
+argnames
 argsort
 ascad
 astrojs
@@ -203,6 +205,7 @@ udevadm
 usbutils
 vals
 vispy
+vmap
 vsintellicode
 webassets
 xlabel

pyproject.toml

@@ -54,6 +54,7 @@ capture = [
 ]
 dev = [
     "flake8",
+    "jax[cuda13]",
     "mypy",
     "psutil",
     "pydantic",

scaaml/stats/cpa/__init__.py

@@ -13,9 +13,21 @@
 # limitations under the License.
 """Correlation power analysis (CPA) module. This module is structured to
 support multiple backends (e.g., a GPU accelerated JAX implementation and a
-NumPy implementation). It currently provides the NumPy version.
+NumPy implementation). If JAX is installed defaults to JAX implementation
+otherwise falls back to the NumPy one.
+
+If a concrete version is needed use:
+```python
+from scaaml.stats.cpa.cpa import CPA
+# from scaaml.stats.cpa.cpa_jax import CPA
+```
 """
-from scaaml.stats.cpa.cpa import CPA  # NumPy based
+try:
+    # JAX based if JAX is installed
+    from scaaml.stats.cpa.cpa_jax import CPA
+except ImportError:
+    # NumPy based default
+    from scaaml.stats.cpa.cpa import CPA  # type: ignore[assignment]
 
 __all__ = [
     "CPA",