Add test for Scan dispatcher with a RandomVariable

Author: rlouf

tests/link/jax/test_scan.py

@@ -38,9 +38,6 @@ def test_sit_sot():
     assert np.allclose(fn(1.0), jax_fn(1.0))
 
 
-@pytest.mark.xfail(
-    reason="Returns correct results but raises exception due to stucture of shared variable."
-)
 def test_nit_sot_shared():
     res, updates = scan(
         fn=lambda: RandomStream(seed=1930, rng_ctor=np.random.RandomState).normal(
@@ -50,9 +47,12 @@ def test_nit_sot_shared():
     )
 
     jax_fn = function((), res, updates=updates, mode="JAX")
-    print(jax_fn())
+    res_jax = jax_fn()
     fn = function((), res, updates=updates)
-    print(fn())
+    res = fn()
+
+    assert res_jax.shape == res.shape
+    assert not np.all(res_jax == res_jax[0])
 
 
 def test_mit_sot():
@@ -80,8 +80,6 @@ def test_mit_sot_2():
     )
     jax_fn = function((), res, updates=updates, mode="JAX")
     fn = function((), res, updates=updates)
-    print(jax_fn())
-    print(fn())
     assert np.allclose(fn(), jax_fn())