Add support for TensorFlow 2.16

.github/workflows/ci.yml

@@ -34,45 +34,49 @@ jobs:
   test:
     needs:
       - static
-    runs-on: ubuntu-latest
     timeout-minutes: 60
     strategy:
       matrix:
         include:
-          - script: remote-test
-            coverage-name: remote
-            tf-version: tensorflow==2.10
+          - script: test
+            coverage-name: gpu
+            tf-version: tensorflow[and-cuda]
+            runs-on: [self-hosted, gpu]
           - script: test
             python-version: "3.11"
             coverage-name: latest
           - script: test
-            tf-version: tensorflow==2.4.4
+            tf-version: tensorflow~=2.6.0
             python-version: "3.8"
             coverage-name: oldest
           - script: test
-            tf-version: tensorflow~=2.6.0
-            python-version: "3.8"
-            coverage-name: tf-2.6
-          - script: remote-docs
-            tf-version: tensorflow==2.10
+            tf-version: tensorflow~=2.9.0
+            coverage-name: tf-2.9
+          - script: test
+            tf-version: tensorflow~=2.13.0
+            coverage-name: tf-2.13
+          - script: docs
+            tf-version: tensorflow[and-cuda]==2.16.1
             python-version: "3.9"
-          - script: remote-examples
-            tf-version: tensorflow==2.10
+            runs-on: [self-hosted, gpu]
+          - script: examples
+            tf-version: tensorflow[and-cuda]==2.16.1
+            runs-on: [self-hosted, gpu]
       fail-fast: false
     env:
       TF_VERSION: ${{ matrix.tf-version || 'tensorflow' }}
-      SSH_KEY: ${{ secrets.SSH_KEY }}
-      SSH_CONFIG: ${{ secrets.SSH_CONFIG }}
-      GH_TOKEN: ${{ secrets.GH_TOKEN }}
+      TF_FORCE_GPU_ALLOW_GROWTH: true
+      GH_TOKEN: ${{ secrets.PUBLIC_GH_TOKEN }}
+    runs-on: ${{ matrix.runs-on || 'ubuntu-latest' }}
     steps:
       - uses: nengo/nengo-bones/actions/setup@main
         with:
           python-version: ${{ matrix.python-version || '3.10' }}
       - uses: nengo/nengo-bones/actions/generate-and-check@main
-      - name: Write secrets to file
+      - name: Install docs requirements
+        if: ${{ contains('docs examples', matrix.script) }}
         run: |
-          mkdir -p ~/.ssh
-          echo '${{ secrets.AZURE_PEM }}' > ~/.ssh/azure.pem
+          micromamba install -y pandoc matplotlib
       - uses: nengo/nengo-bones/actions/run-script@main
         with:
           name: ${{ matrix.script }}

.nengobones.yml

@@ -17,14 +17,15 @@ manifest_in: {}
 
 setup_py:
   install_req:
+    - anyio<4  # not compatible with older tensorflow versions
     - packaging>=20.9
     - scipy>=1.0.0
-    - tensorflow>=2.4.4
+    - tensorflow>=2.6.0
   tests_req:
     - pytest>=6.1.0
     - pytest-rng>=1.0.0
   docs_req:
-    - matplotlib>=3.0.2,<3.4.3
+    - matplotlib>=3.8.4
     - jupyter>=1.0.0
     - seaborn>=0.9.0
     - sphinx>=1.8
@@ -66,39 +67,15 @@ docs_conf_py:
 ci_scripts:
   - template: static
   - template: docs
+    pre_commands:
+      # We run this ahead of time, otherwise the download progress bar causes
+      # problems in the notebook rendering
+      - python -c "import tensorflow as tf; tf.keras.datasets.mnist.load_data()"
   - template: examples
   - template: test
     coverage: true
     pip_install:
       - $TF_VERSION
-  - template: remote-script
-    remote_script: test
-    output_name: remote-test
-    host: azure
-    azure_name: nengo-dl
-    azure_group: nengo-ci
-    coverage: true
-    remote_vars:
-      TF_FORCE_GPU_ALLOW_GROWTH: "true"
-      TF_VERSION: $TF_VERSION
-    remote_setup:
-      - micromamba install -y "$TF_VERSION" cudnn=8.4
-  - template: remote-script
-    remote_script: docs
-    output_name: remote-docs
-    host: azure-docs
-    azure_name: nengo-dl-docs
-    azure_group: nengo-ci
-    remote_setup:
-      - micromamba install -y "$TF_VERSION" cudnn=8.4
-  - template: remote-script
-    remote_script: examples
-    output_name: remote-examples
-    host: azure-examples
-    azure_name: nengo-dl-examples
-    azure_group: nengo-ci
-    remote_setup:
-      - micromamba install -y "$TF_VERSION" cudnn=8.4
   - template: deploy
     wheel: true
 
@@ -109,6 +86,6 @@ pyproject_toml: {}
 version_py:
   type: semver
   major: 0
-  minor: 7
-  patch: 1
+  minor: 8
+  patch: 0
   release: false

CHANGES.rst

@@ -19,10 +19,10 @@ Release history
    - Removed
    - Fixed
 
-0.7.1 (unreleased)
+0.8.0 (unreleased)
 ==================
 
-*Compatible with TensorFlow 2.4 - 2.13*
+*Compatible with TensorFlow 2.6 - 2.16*
 
 0.7.0 (July 20, 2023)
 =====================

docs/basic-usage.rst

@@ -11,13 +11,14 @@ a new LMU layer:
 
 .. testcode::
 
+   import keras
    import keras_lmu
 
    lmu_layer = keras_lmu.LMU(
        memory_d=1,
        order=256,
        theta=784,
-       hidden_cell=tf.keras.layers.SimpleRNNCell(units=10),
+       hidden_cell=keras.layers.SimpleRNNCell(units=10),
    )
 
 Note that the values used above for ``memory_d``, ``order``,

docs/examples/psMNIST.ipynb

@@ -42,6 +42,7 @@
    "source": [
     "%matplotlib inline\n",
     "\n",
+    "import keras\n",
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
     "from IPython.display import Image, display\n",
@@ -261,22 +262,22 @@
     "    memory_d=1,\n",
     "    order=256,\n",
     "    theta=n_pixels,\n",
-    "    hidden_cell=tf.keras.layers.SimpleRNNCell(212),\n",
+    "    hidden_cell=keras.layers.SimpleRNNCell(212),\n",
     "    hidden_to_memory=False,\n",
     "    memory_to_memory=False,\n",
     "    input_to_hidden=True,\n",
     "    kernel_initializer=\"ones\",\n",
     ")\n",
     "\n",
     "# TensorFlow layer definition\n",
-    "inputs = tf.keras.Input((n_pixels, 1))\n",
+    "inputs = keras.Input((n_pixels, 1))\n",
     "lmus = lmu_layer(inputs)\n",
-    "outputs = tf.keras.layers.Dense(10)(lmus)\n",
+    "outputs = keras.layers.Dense(10)(lmus)\n",
     "\n",
     "# TensorFlow model definition\n",
-    "model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
+    "model = keras.Model(inputs=inputs, outputs=outputs)\n",
     "model.compile(\n",
-    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
     "    optimizer=\"adam\",\n",
     "    metrics=[\"accuracy\"],\n",
     ")\n",
@@ -316,10 +317,13 @@
     "batch_size = 100\n",
     "epochs = 10\n",
     "\n",
-    "saved_weights_fname = \"./psMNIST-weights.hdf5\"\n",
+    "saved_model_fname = \"./psMNIST.keras\"\n",
     "callbacks = [\n",
-    "    tf.keras.callbacks.ModelCheckpoint(\n",
-    "        filepath=saved_weights_fname, monitor=\"val_loss\", verbose=1, save_best_only=True\n",
+    "    keras.callbacks.ModelCheckpoint(\n",
+    "        filepath=saved_model_fname,\n",
+    "        monitor=\"val_accuracy\",\n",
+    "        verbose=1,\n",
+    "        save_best_only=True,\n",
     "    ),\n",
     "]\n",
     "\n",
@@ -393,8 +397,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model.load_weights(saved_weights_fname)\n",
-    "accuracy = model.evaluate(X_test, Y_test)[1] * 100\n",
+    "model.load_weights(saved_model_fname)\n",
+    "accuracy = model.evaluate(X_test, Y_test, verbose=0)[1] * 100\n",
     "print(f\"Test accuracy: {round(accuracy, 2):0.2f}%\")"
    ]
   },