Support for TensorFlow 2.4

CHANGES.rst

@@ -26,6 +26,7 @@ Release history
 - Added support for `KerasSpiking <https://www.nengo.ai/keras-spiking/>`_ layers in
   the Converter. (`#182`_)
 - Added support for ``tf.keras.layers.TimeDistributed`` in the Converter. (`#182`_)
+- Added support for TensorFlow 2.4. (`#185`_)
 
 **Changed**
 
@@ -47,6 +48,7 @@ Release history
 .. _#173: https://github.com/nengo/nengo-dl/pull/173
 .. _#181: https://github.com/nengo/nengo-dl/pull/181
 .. _#182: https://github.com/nengo/nengo-dl/pull/182
+.. _#185: https://github.com/nengo/nengo-dl/pull/185
 
 3.3.0 (August 14, 2020)
 -----------------------

nengo_dl/converter.py

@@ -391,9 +391,7 @@ def _get_key(self, key):
                 # get output tensor
                 key = key.output
 
-            key = tuple(
-                x.ref() if isinstance(x, tf.Tensor) else x for x in tf.nest.flatten(key)
-            )
+            key = tuple(x.ref() if tf.is_tensor(x) else x for x in tf.nest.flatten(key))
 
             return key
 

nengo_dl/neuron_builders.py

@@ -11,7 +11,7 @@
 from nengo.neurons import RectifiedLinear, SpikingRectifiedLinear, Sigmoid, LIF, LIFRate
 import numpy as np
 import tensorflow as tf
-from tensorflow.python.keras.utils import tf_utils
+from tensorflow.python.framework import smart_cond
 
 from nengo_dl import compat, utils
 from nengo_dl.builder import Builder, OpBuilder
@@ -188,7 +188,7 @@ def build_step(self, signals, **step_kwargs):
             out = step_output
         else:
             out = tf.nest.flatten(
-                tf_utils.smart_cond(
+                smart_cond.smart_cond(
                     self.config.training,
                     true_fn=lambda: (self.training_step(J, signals.dt, **state),)
                     + tuple(state.values()),

nengo_dl/simulator.py

@@ -1815,7 +1815,9 @@ def _generate_inputs(self, data=None, n_steps=None):
         if data is None:
             data = {}
 
-        if not isinstance(data, (list, tuple, dict, np.ndarray, tf.Tensor)):
+        if not isinstance(data, (list, tuple, dict, np.ndarray)) and not tf.is_tensor(
+            data
+        ):
             # data is some kind of generator, so we don't try to modify it (too many
             # different types of generators this could be)
             if n_steps is not None:

nengo_dl/tensor_node.py

@@ -41,7 +41,7 @@ def validate_output(output, minibatch_size=None, output_d=None, dtype=None):
         Expected dtype of the function output.
     """
 
-    if not isinstance(output, (tf.Tensor, tf.TensorSpec)):
+    if not isinstance(output, tf.TensorSpec) and not tf.is_tensor(output):
         raise ValidationError(
             "TensorNode function must return a Tensor (got %s)" % type(output),
             attr="tensor_func",

nengo_dl/tests/test_simulator.py

@@ -837,7 +837,16 @@ def test_tensorboard(Simulator, tmpdir):
 
 @pytest.mark.parametrize("mode", ("predict", "train"))
 @pytest.mark.training
-def test_profile(Simulator, mode, tmpdir):
+def test_profile(Simulator, mode, tmpdir, pytestconfig):
+    if (
+        pytestconfig.getoption("--graph-mode")
+        and version.parse(tf.__version__) >= version.parse("2.4.0rc0")
+        and mode == "predict"
+    ):
+        pytest.skip(
+            "TensorFlow bug, see https://github.com/tensorflow/tensorflow/issues/44563"
+        )
+
     net, a, p = dummies.linear_net()
 
     with Simulator(net) as sim:

nengo_dl/tests/test_tensor_node.py

@@ -373,7 +373,7 @@ def __init__(self, expected):
             self.expected = expected
 
         def call(self, inputs, training=None):
-            tf.assert_equal(training, self.expected)
+            tf.assert_equal(tf.cast(training, tf.bool), self.expected)
             return tf.reshape(inputs, (1, 1))
 
     with nengo.Network() as net: