Support for TensorFlow 2.4

.nengobones.yml

@@ -71,6 +71,8 @@ setup_cfg:
       tests/test_probe.py::test_multirun:
         simulation times may not line up exactly if unroll_simulation != 1, see
         tests/test_nengo_tests.py::test_multirun
+      tests/test_learning_rules.py::test_rls_*:
+        RLS learning rule not implemented
     allclose_tolerances:
       - tests/test_synapses.py::test_lowpass atol=5e-7
       - tests/test_synapses.py::test_triangle atol=5e-7
@@ -137,6 +139,8 @@ docs_conf_py:
 
 ci_scripts:
   - template: static
+    pre_commands:
+      - export SHELLCHECK_OPTS="-e SC2086"
   - template: test
     pip_install:
       - $NUMPY_VERSION
@@ -158,46 +162,87 @@ ci_scripts:
       # TODO: swap to pypi release after 0.2
       - git+https://github.com/nengo/keras-spiking
     nengo_tests: true
+  - template: test
+    output_name: test-performance
+    coverage: true
+    pip_install:
+      - $NUMPY_VERSION
+      - $TF_VERSION
+      - $NENGO_VERSION
+      - $SCIPY_VERSION
+      - git+https://github.com/drasmuss/spaun2.0.git
+    nengo_tests: true
+    pre_commands:
+      - pytest $TEST_ARGS nengo_dl/tests/test_benchmarks.py::test_performance --performance -v --durations 20 --color=yes
   - template: examples
     pip_install:
       - $TF_VERSION
+      - $NUMPY_VERSION
   - template: docs
     pip_install:
       - $TF_VERSION
       - $NUMPY_VERSION
-  - template: gpu
+  - template: remote-script
+    remote_script: test-performance
+    output_name: remote-test
     host: azure
     travis_var_key: fcdbeb54a1c4
     coverage: true
+    azure_name: nengo-dl
+    azure_group: nengo-ci
+    remote_vars:
+      TF_FORCE_GPU_ALLOW_GROWTH: "true"
+      NUMPY_VERSION: $NUMPY_VERSION
+      TF_VERSION: $TF_VERSION
+      NENGO_VERSION: $NENGO_VERSION
+      SCIPY_VERSION: $SCIPY_VERSION
+    remote_setup:
+      - conda install -y cudatoolkit=10.1 cudnn
   - template: remote-script
     remote_script: docs
     output_name: remote-docs
     host: azure-docs
     travis_var_key: fcdbeb54a1c4
+    azure_name: nengo-dl-docs
+    azure_group: nengo-ci
+    remote_vars:
+      NUMPY_VERSION: $NUMPY_VERSION
+      TF_VERSION: $TF_VERSION
+      NENGO_VERSION: $NENGO_VERSION
+      SCIPY_VERSION: $SCIPY_VERSION
+    remote_setup:
+      - conda install -y cudatoolkit=10.1 cudnn
   - template: remote-script
     remote_script: examples
     output_name: remote-examples
     host: azure-examples
     travis_var_key: fcdbeb54a1c4
+    azure_name: nengo-dl-examples
+    azure_group: nengo-ci
+    remote_vars:
+      NUMPY_VERSION: $NUMPY_VERSION
+      TF_VERSION: $TF_VERSION
+      NENGO_VERSION: $NENGO_VERSION
+      SCIPY_VERSION: $SCIPY_VERSION
+    remote_setup:
+      - conda install -y cudatoolkit=10.1 cudnn
   - template: deploy
 
 travis_yml:
   python: 3.6
   global_vars:
     NUMPY_VERSION: numpy>=1.16.0
-    TF_VERSION: tensorflow-cpu
+    TF_VERSION: tensorflow
     NENGO_VERSION: nengo[tests]
     SCIPY_VERSION: noop
   jobs:
     - stage: basic
       script: static
-    - cache: false
-    - script: gpu
-      test_args: --device /gpu:0
+    - script: test
+    - script: remote-test
       env:
         NENGO_VERSION: git+https://github.com/nengo/nengo.git#egg=nengo[tests]
-        TF_VERSION: tensorflow-gpu
-        GPU_NUM: 0
+        TF_VERSION: tensorflow<2.3.0
     - script: test-coverage
       env:
         NENGO_VERSION: git+https://github.com/nengo/nengo.git#egg=nengo[tests]
@@ -214,16 +259,10 @@ travis_yml:
       test_args: --graph-mode
     - stage: advanced
       script: remote-docs
-      env:
-        TF_VERSION: tensorflow-gpu
-        GPU_NUM: 0
     - script: remote-examples
-      env:
-        TF_VERSION: tensorflow-gpu
-        GPU_NUM: 0
     - test_args: --unroll-simulation 5 --dtype float64 --simulator-only
     - test_args: --inference-only --simulator-only
-  pypi_user: drasmuss
+  pypi_user: __token__
   deploy_dists:
     - sdist
   slack_notifications: "u7LEiuoLupKwd/p4C3MXaojqvkBvZKlSw6UaBr2C6ZJINRD2N/34bHtv6Qc1wi\

.travis.yml

@@ -24,7 +24,7 @@ env:
     - BRANCH_NAME="${TRAVIS_PULL_REQUEST_BRANCH:-$TRAVIS_BRANCH}"
     - PIP_USE_FEATURE="2020-resolver"
     - NUMPY_VERSION="numpy>=1.16.0"
-    - TF_VERSION="tensorflow-cpu"
+    - TF_VERSION="tensorflow"
     - NENGO_VERSION="nengo[tests]"
     - SCIPY_VERSION="noop"
 
@@ -35,14 +35,13 @@ jobs:
       SCRIPT="static"
     stage: basic
   -
-    cache: False
+    env:
+      SCRIPT="test"
   -
     env:
       NENGO_VERSION="git+https://github.com/nengo/nengo.git#egg=nengo[tests]"
-      TF_VERSION="tensorflow-gpu"
-      GPU_NUM="0"
-      SCRIPT="gpu"
-      TEST_ARGS="--device /gpu:0"
+      TF_VERSION="tensorflow<2.3.0"
+      SCRIPT="remote-test"
   -
     env:
       NENGO_VERSION="git+https://github.com/nengo/nengo.git#egg=nengo[tests]"
@@ -63,14 +62,10 @@ jobs:
       TEST_ARGS="--graph-mode"
   -
     env:
-      TF_VERSION="tensorflow-gpu"
-      GPU_NUM="0"
       SCRIPT="remote-docs"
     stage: advanced
   -
     env:
-      TF_VERSION="tensorflow-gpu"
-      GPU_NUM="0"
       SCRIPT="remote-examples"
   -
     env:
@@ -85,15 +80,15 @@ jobs:
     deploy:
       - provider: pypi
         server: https://test.pypi.org/legacy/
-        user: drasmuss
-        password: $PYPI_TOKEN
+        user: __token__
+        password: $PYPI_TEST_TOKEN
         distributions: "sdist "
         on:
           all_branches: true
           tags: false
           condition: $TRAVIS_BRANCH =~ ^release-candidate-*
       - provider: pypi
-        user: drasmuss
+        user: __token__
         password: $PYPI_TOKEN
         distributions: "sdist "
         on:

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_graph.py

@@ -110,6 +110,14 @@ def __init__(
             )
         ]
 
+        # check for unsupported operators
+        for op in operators:
+            if type(op) not in builder.Builder.builders:
+                raise BuildError(
+                    "No registered builder for operators of type %s; "
+                    "consider registering a custom builder" % type(op)
+                )
+
         # mark trainable signals
         self.mark_signals()
 

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: