Case where grad shape is wrong

I have come across a case where grad of a scalar with respect to another scalar gives a vector. I have noticed at least one place where if you change the code you get a scalar (indicated by the fail clause).

• Check for duplicate issues.
• Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

from jax import numpy as jnp, vmap, value_and_grad

FAIL = True


def f2(y, z):
    v1 = z
    v2 = jnp.sum(y) + z

    if FAIL:
        return jnp.logaddexp(v1, v2)
    else:
        return v1 + v2


def f1(y, z):
    v = vmap(lambda _y: f2(_y, z))(y)
    return jnp.sum(v)


if __name__ == '__main__':
    y = jnp.ones((3, 2))
    f = lambda z: f1(y, z)
    z = 0.1
    val, g = value_and_grad(f)(z)
    print(val.shape, g.shape)
    assert val.shape == ()
    assert g.shape == ()

jax.__version__ = 0.2.9
jaxlib.__version__ = 0.1.61

[Joshuaalbert]

#6470 fixed it in both the MVCE and the larger algorithm.

[mattjj]

I think I can get you a better fix sometime today.