Add test with fixed number of steps

Author: rlouf

tests/link/jax/test_scan.py

@@ -2,6 +2,7 @@
 import pytest
 from packaging.version import parse as version_parse
 
+import aesara
 import aesara.tensor as at
 from aesara.configdefaults import config
 from aesara.graph.fg import FunctionGraph
@@ -14,6 +15,28 @@
 jax = pytest.importorskip("jax")
 
 
+def test_jax_fixed_length():
+    """
+
+    In this test we make sure that we can compile a scan loop
+    with a fixed number of steps.
+
+    """
+    A = at.vector("A")
+
+    result, _ = aesara.scan(
+        fn=lambda prior_result, A: prior_result * A,
+        outputs_info=at.ones_like(A),
+        non_sequences=A,
+        n_steps=4,
+    )
+
+    final_result = result
+
+    out_fg = FunctionGraph([A], [final_result])
+    compare_jax_and_py(out_fg, [np.array([2.0])])
+
+
 @pytest.mark.xfail(
     version_parse(jax.__version__) >= version_parse("0.2.12"),
     reason="Omnistaging cannot be disabled",