add test to jax output shape

Author: 1b15

tests/test_jax.py

@@ -104,6 +104,9 @@ def test_network(self):
 
         model_jax.run(chunkwise=True, bold=True, append_outputs=True, chunksize=20000)
 
+        # assert same output shape
+        self.assertTrue(model.exc.shape == model_jax.exc.shape)
+
         # jit changes the exact numerics of outputs
         self.assertTrue(np.allclose(model.BOLD.BOLD, model_jax.BOLD.BOLD, rtol=1e-3))