Add test for scan with multiple None outputs

Author: rlouf

tests/link/jax/test_scan.py

@@ -356,3 +356,29 @@ def input_step_fn(y_tm1, y_tm3, a):
 
     test_input_vals = [np.array(10.0).astype(config.floatX)]
     compare_jax_and_py(out_fg, test_input_vals)
+
+
+def test_scan_multiple_none_output():
+    A = at.dvector("A")
+
+    def power_step(prior_result, x):
+        return prior_result * x, prior_result * x * x, prior_result * x * x * x
+
+    result, _ = scan(
+        power_step,
+        non_sequences=[A],
+        outputs_info=[at.ones_like(A), None, None],
+        n_steps=3,
+    )
+
+    FunctionGraph([A], result)
+    test_input_vals = (np.array([1.0, 2.0]),)
+
+    jax_fn = function((A,), result, mode="JAX")
+    jax_res = jax_fn(*test_input_vals)
+
+    fn = function((A,), result)
+    res = fn(*test_input_vals)
+
+    for output_jax, output in zip(jax_res, res):
+        assert np.allclose(jax_res, res)