From 7264265a280db0e265c8dc20004af6d7fb02d264 Mon Sep 17 00:00:00 2001
From: Daniel Rasmussen <dhrsmss@gmail.com>
Date: Mon, 22 Apr 2024 17:05:29 -0300
Subject: [PATCH] Update for Keras 3.3

---
 .nengobones.yml             |  2 +-
 docs/basic-usage.rst        |  9 +++------
 docs/examples/psMNIST.ipynb |  6 ++----
 keras_lmu/layers.py         | 11 ++++++++---
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/.nengobones.yml b/.nengobones.yml
index 92730e7..0be4120 100644
--- a/.nengobones.yml
+++ b/.nengobones.yml
@@ -70,7 +70,7 @@ ci_scripts:
     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()"
+      - python -c "import keras; keras.datasets.mnist.load_data()"
   - template: examples
   - template: test
     coverage: true
diff --git a/docs/basic-usage.rst b/docs/basic-usage.rst
index 62c7dcb..968fb64 100644
--- a/docs/basic-usage.rst
+++ b/docs/basic-usage.rst
@@ -41,14 +41,11 @@ a 10-dimensional input and a 20-dimensional output.
 
 .. testcode::
 
-   from tensorflow.keras import Input, Model
-   from tensorflow.keras.layers import Dense
-
-   inputs = Input((None, 10))
+   inputs = keras.Input((None, 10))
    lmus = lmu_layer(inputs)
-   outputs = Dense(20)(lmus)
+   outputs = keras.layers.Dense(20)(lmus)
 
-   model = Model(inputs=inputs, outputs=outputs)
+   model = keras.Model(inputs=inputs, outputs=outputs)
 
 
 Other parameters
diff --git a/docs/examples/psMNIST.ipynb b/docs/examples/psMNIST.ipynb
index 9167c77..f241820 100644
--- a/docs/examples/psMNIST.ipynb
+++ b/docs/examples/psMNIST.ipynb
@@ -82,9 +82,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "We now obtain the standard MNIST dataset of handwritten digits from `tf.keras.datasets`."
-   ]
+   "source": "We now obtain the standard MNIST dataset of handwritten digits from `keras.datasets`."
   },
   {
    "cell_type": "code",
@@ -95,7 +93,7 @@
     "(train_images, train_labels), (\n",
     "    test_images,\n",
     "    test_labels,\n",
-    ") = tf.keras.datasets.mnist.load_data()"
+    ") = keras.datasets.mnist.load_data()"
    ]
   },
   {
diff --git a/keras_lmu/layers.py b/keras_lmu/layers.py
index 22b2003..982a4f2 100644
--- a/keras_lmu/layers.py
+++ b/keras_lmu/layers.py
@@ -29,8 +29,13 @@
 else:
     from keras.layers import Layer as BaseRandomLayer
 
+if tf_version < version.parse("2.16.0rc0"):
+    from tensorflow.keras.utils import register_keras_serializable
+else:
+    from keras.saving import register_keras_serializable
+
 
-@tf.keras.utils.register_keras_serializable("keras-lmu")
+@register_keras_serializable("keras-lmu")
 class LMUCell(
     DropoutRNNCellMixin, BaseRandomLayer
 ):  # pylint: disable=too-many-ancestors
@@ -524,7 +529,7 @@ def from_config(cls, config):
         return super().from_config(config)
 
 
-@tf.keras.utils.register_keras_serializable("keras-lmu")
+@register_keras_serializable("keras-lmu")
 class LMU(keras.layers.Layer):  # pylint: disable=too-many-ancestors,abstract-method
     """
     A layer of trainable low-dimensional delay systems.
@@ -792,7 +797,7 @@ def from_config(cls, config):
         return super().from_config(config)
 
 
-@tf.keras.utils.register_keras_serializable("keras-lmu")
+@register_keras_serializable("keras-lmu")
 class LMUFeedforward(
     keras.layers.Layer
 ):  # pylint: disable=too-many-ancestors,abstract-method