Add tests for shared updates in JAX and NUMBA backends

aesara/link/basic.py

@@ -641,7 +641,8 @@ class JITLinker(PerformLinker):
             The JITed function that performs the computations.
 
         """
-        output_nodes = [o.owner for o in self.fgraph.outputs]
+        # This is a bit hackish, but we only return one of the output nodes
+        output_nodes = [o.owner for o in self.fgraph.outputs if o.owner is not None][:1]
 
         converted_fgraph = self.fgraph_convert(
             self.fgraph,
@@ -678,8 +679,7 @@ class JITLinker(PerformLinker):
 
         thunks.append(thunk)
 
-        # This is a bit hackish, but we only return one of the output nodes
-        return thunks, output_nodes[:1], fgraph_jit
+        return thunks, output_nodes, fgraph_jit
 
     def make_all(self, input_storage=None, output_storage=None, storage_map=None):
         fgraph = self.fgraph

aesara/link/utils.py

@@ -243,7 +243,7 @@ def gc_helper(node_list: List[Apply]):
     -------
     2-tuple
         FIRST, the set of Variable instances which are computed by node_list,
-        and SECOND a dictionary that maps each Variable instance to a the last
+        and SECOND a dictionary that maps each Variable instance to the last
         node to use Variable as an input.
 
     Extended Summary

tests/link/jax/test_basic.py

@@ -182,6 +182,22 @@ def test_shared():
     np.testing.assert_allclose(jax_res, new_a_value * 2)
 
 
+def test_shared_updates():
+    a = shared(0)
+
+    aesara_jax_fn = function([], a, updates={a: a + 1}, mode="JAX")
+    res1, res2 = aesara_jax_fn(), aesara_jax_fn()
+    assert res1 == 0
+    assert res2 == 1
+    assert a.get_value() == 2
+
+    a.set_value(5)
+    res1, res2 = aesara_jax_fn(), aesara_jax_fn()
+    assert res1 == 5
+    assert res2 == 6
+    assert a.get_value() == 7
+
+
 def test_jax_ifelse():
 
     true_vals = np.r_[1, 2, 3]

tests/link/numba/test_basic.py

@@ -837,6 +837,22 @@ def test_shared():
     np.testing.assert_allclose(numba_res, new_a_value * 2)
 
 
+def test_shared_updates():
+    a = shared(0)
+
+    aesara_numba_fn = function([], a, updates={a: a + 1}, mode="NUMBA")
+    res1, res2 = aesara_numba_fn(), aesara_numba_fn()
+    assert res1 == 0
+    assert res2 == 1
+    assert a.get_value() == 2
+
+    a.set_value(5)
+    res1, res2 = aesara_numba_fn(), aesara_numba_fn()
+    assert res1 == 5
+    assert res2 == 6
+    assert a.get_value() == 7
+
+
 # We were seeing some weird results in CI where the following two almost
 # sign-swapped results were being return from Numba and Python, respectively.
 # The issue might be related to https://github.com/numba/numba/issues/4519.