Revert "fix + test"

This reverts commit 4cd79c2.

equinox/_ad.py

@@ -439,10 +439,7 @@ def _fun(_diff_x):
         else:
             jacobian = jax.jacfwd
         dynamic_out, (static_out, aux) = jacobian(_fun, has_aux=True)(diff_x)
-        if static_out is None:
-            out = dynamic_out
-        else:
-            out = combine(dynamic_out, static_out)
+        out = combine(dynamic_out, static_out)
         if self.has_aux:
             return out, aux
         else:

tests/test_ad.py

@@ -483,41 +483,6 @@ def f_bwd(res, g):
     assert tree_allclose(jax_hess, eqx_hess)
 
 
-def test_pytree_jacfwd():
-    class NeuralNetwork(eqx.Module):
-        layers: list
-        extra_bias: jax.Array
-
-        def __init__(self, key):
-            key1, key2, key3 = jax.random.split(key, 3)
-            self.layers = [
-                eqx.nn.Linear(2, 8, key=key1),
-                eqx.nn.Linear(8, 8, key=key2),
-                eqx.nn.Linear(8, 2, key=key3),
-            ]
-            self.extra_bias = jax.numpy.ones(2)
-
-        def __call__(self, x):
-            for layer in self.layers[:-1]:
-                x = jax.nn.relu(layer(x))
-            return self.layers[-1](x) + self.extra_bias
-
-    def loss(model, x, y):
-        pred_y = jax.vmap(model)(x)
-        return jax.numpy.mean((y - pred_y) ** 2)
-
-    x_key, y_key, model_key = jax.random.split(jax.random.PRNGKey(0), 3)
-    x = jax.random.normal(x_key, (3, 2))
-    y = jax.random.normal(y_key, (3, 2))
-    model = NeuralNetwork(model_key)
-    assert tree_allclose(
-        eqx.filter_grad(loss)(model, x, y), eqx.filter_jacfwd(loss)(model, x, y)
-    )
-    assert tree_allclose(
-        eqx.filter_grad(loss)(model, x, y), eqx.filter_jacrev(loss)(model, x, y)
-    )
-
-
 def test_filter_custom_jvp_symbolic_zero():
     @eqx.filter_custom_jvp
     def f(x, y):