WIP, ENH: Array API for GFDLClassifier

.github/workflows/array-api.yml

@@ -0,0 +1,34 @@
+name: Array API CI
+
+on:
+  push:
+    branches:
+      - maintenance/**
+  pull_request:
+    branches:
+      - main
+      - maintenance/**
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.14"]
+
+    steps:
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install -v .
+        python -m pip install -v --group test
+        python -m pip install torch array-api-strict array-api-compat
+    - name: test
+      run: |
+        cd $RUNNER_TEMP
+        SCIPY_ARRAY_API=1 python -m pytest --pyargs gfdl --cov=gfdl --cov-report=term-missing

src/gfdl/activations.py

@@ -2,8 +2,7 @@
 Activation functions for Gradient Free Deep Learning estimators.
 """
 
-import numpy as np
-import scipy
+from sklearn.utils._array_api import get_namespace_and_device
 
 
 def relu(z):
@@ -17,10 +16,12 @@ def relu(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The output array with only positive values.
     """
-    return np.maximum(0, z)
+    xp, _, device = get_namespace_and_device(z)
+    zero = xp.asarray(0, dtype=z.dtype, device=device)
+    return xp.maximum(zero, z)
 
 
 def tanh(z):
@@ -34,14 +35,15 @@ def tanh(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The output array with hyperbolic tangent values.
 
     See Also
     --------
     numpy.tanh : The hyperbolic tangent function.
     """
-    return np.tanh(z)
+    xp, _, _device = get_namespace_and_device(z)
+    return xp.tanh(z)
 
 
 def sigmoid(z):
@@ -55,14 +57,15 @@ def sigmoid(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The output array with the function values.
 
     See Also
     --------
     scipy.special.expit : The logistic sigmoid function.
     """
-    return scipy.special.expit(z)
+    xp, _, _ = get_namespace_and_device(z)
+    return xp.exp(-xp.logaddexp(xp.zeros_like(z), -z))
 
 
 def identity(z):
@@ -76,7 +79,7 @@ def identity(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The input array is returned unchanged.
     """
     return z
@@ -93,14 +96,17 @@ def softmax(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The output array with the function values.
 
     See Also
     --------
     scipy.special.softmax : The softmax function.
     """
-    return scipy.special.softmax(z, axis=-1)
+    xp, _, _ = get_namespace_and_device(z)
+    tmp = z - xp.max(z, axis=1, keepdims=True)
+    exp_tmp = xp.exp(tmp)
+    return exp_tmp / xp.sum(exp_tmp, axis=1, keepdims=True)
 
 
 def softmin(z):
@@ -116,7 +122,7 @@ def softmin(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The output array with the function values.
 
     See Also
@@ -137,14 +143,15 @@ def log_sigmoid(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The output array with the function values.
 
     See Also
     --------
     scipy.special.log_expit : The logistic sigmoid function.
     """
-    return scipy.special.log_expit(z)
+    xp, _, _ = get_namespace_and_device(z)
+    return -xp.logaddexp(xp.zeros_like(z), -z)
 
 
 def log_softmax(z):
@@ -158,14 +165,18 @@ def log_softmax(z):
 
     Returns
     -------
-    numpy.ndarray
+    array
         The output array with the function values.
 
     See Also
     --------
     scipy.special.log_softmax : The logarithm of softmax function.
     """
-    return scipy.special.log_softmax(z, axis=-1)
+    xp, _, _ = get_namespace_and_device(z)
+    shifted = z - xp.max(z, axis=1, keepdims=True)
+    return shifted - xp.log(
+        xp.sum(xp.exp(shifted), axis=1, keepdims=True)
+    )
 
 
 ACTIVATIONS = {