iterate over out axes when vmapped twice

Author: jpbrodrick89

tests/test_endtoend.py

@@ -221,17 +221,22 @@ def f(x, y):
                 extra_dim if ax is not None else None for ax in axes
             )
 
-            f_vmappedtwice = jax.vmap(f_vmapped, in_axes=additional_axes)
-            raw_vmappedtwice = jax.vmap(raw_vmapped, in_axes=additional_axes)
-
-            if use_jit:
-                f_vmappedtwice = jax.jit(f_vmappedtwice)
-                raw_vmappedtwice = jax.jit(raw_vmappedtwice)
-
-            result = f_vmappedtwice(x, y)
-            result_raw = raw_vmappedtwice(x, y)
-
-            _assert_pytree_isequal(result, result_raw)
+            for out_axis in [0, 1, -1]:
+                f_vmappedtwice = jax.vmap(
+                    f_vmapped, in_axes=additional_axes, out_axes=out_axis
+                )
+                raw_vmappedtwice = jax.vmap(
+                    raw_vmapped, in_axes=additional_axes, out_axes=out_axis
+                )
+
+                if use_jit:
+                    f_vmappedtwice = jax.jit(f_vmappedtwice)
+                    raw_vmappedtwice = jax.jit(raw_vmappedtwice)
+
+                result = f_vmappedtwice(x, y)
+                result_raw = raw_vmappedtwice(x, y)
+
+                _assert_pytree_isequal(result, result_raw)
 
 
 @pytest.mark.parametrize("use_jit", [True, False])