<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto"><code class="notranslate">jax.numpy.prod</code> does not work with tuples but <code class="notranslate">numpy.prod</code> works just fine. See MWEs below</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import numpy
print(numpy.prod((1,128)))
# 128
print(numpy.prod((128,)))
#128"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">numpy</span>
<span class="pl-en">print</span>(<span class="pl-s1">numpy</span>.<span class="pl-en">prod</span>((<span class="pl-c1">1</span>,<span class="pl-c1">128</span>)))
<span class="pl-c"># 128</span>
<span class="pl-en">print</span>(<span class="pl-s1">numpy</span>.<span class="pl-en">prod</span>((<span class="pl-c1">128</span>,)))
<span class="pl-c">#128</span></pre></div>
<p dir="auto">The same code crashes for <code class="notranslate">jax.numpy</code></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
print(jax.numpy.prod((1,128)))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-en">print</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-en">prod</span>((<span class="pl-c1">1</span>,<span class="pl-c1">128</span>)))</pre></div>
<p dir="auto">Error mesage:<br>
<code class="notranslate">TypeError: prod requires ndarray or scalar arguments, got &lt;class 'tuple'&gt; at position 0.</code></p>
<p dir="auto">I have come across a case where <code class="notranslate">grad</code> 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).</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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 == ()
"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>, <span class="pl-s1">vmap</span>, <span class="pl-s1">value_and_grad</span>

<span class="pl-v">FAIL</span> <span class="pl-c1">=</span> <span class="pl-c1">True</span>


<span class="pl-k">def</span> <span class="pl-en">f2</span>(<span class="pl-s1">y</span>, <span class="pl-s1">z</span>):
    <span class="pl-s1">v1</span> <span class="pl-c1">=</span> <span class="pl-s1">z</span>
    <span class="pl-s1">v2</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">y</span>) <span class="pl-c1">+</span> <span class="pl-s1">z</span>

    <span class="pl-k">if</span> <span class="pl-v">FAIL</span>:
        <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">logaddexp</span>(<span class="pl-s1">v1</span>, <span class="pl-s1">v2</span>)
    <span class="pl-k">else</span>:
        <span class="pl-k">return</span> <span class="pl-s1">v1</span> <span class="pl-c1">+</span> <span class="pl-s1">v2</span>


<span class="pl-k">def</span> <span class="pl-en">f1</span>(<span class="pl-s1">y</span>, <span class="pl-s1">z</span>):
    <span class="pl-s1">v</span> <span class="pl-c1">=</span> <span class="pl-en">vmap</span>(<span class="pl-k">lambda</span> <span class="pl-s1">_y</span>: <span class="pl-en">f2</span>(<span class="pl-s1">_y</span>, <span class="pl-s1">z</span>))(<span class="pl-s1">y</span>)
    <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">v</span>)


<span class="pl-k">if</span> <span class="pl-s1">__name__</span> <span class="pl-c1">==</span> <span class="pl-s">'__main__'</span>:
    <span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>((<span class="pl-c1">3</span>, <span class="pl-c1">2</span>))
    <span class="pl-s1">f</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">z</span>: <span class="pl-en">f1</span>(<span class="pl-s1">y</span>, <span class="pl-s1">z</span>)
    <span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-c1">0.1</span>
    <span class="pl-s1">val</span>, <span class="pl-s1">g</span> <span class="pl-c1">=</span> <span class="pl-en">value_and_grad</span>(<span class="pl-s1">f</span>)(<span class="pl-s1">z</span>)
    <span class="pl-en">print</span>(<span class="pl-s1">val</span>.<span class="pl-s1">shape</span>, <span class="pl-s1">g</span>.<span class="pl-s1">shape</span>)
    <span class="pl-k">assert</span> <span class="pl-s1">val</span>.<span class="pl-s1">shape</span> <span class="pl-c1">==</span> ()
    <span class="pl-k">assert</span> <span class="pl-s1">g</span>.<span class="pl-s1">shape</span> <span class="pl-c1">==</span> ()</pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="jax.__version__ = 0.2.9
jaxlib.__version__ = 0.1.61"><pre class="notranslate"><span class="pl-s1">jax</span>.<span class="pl-s1">__version__</span> <span class="pl-c1">=</span> <span class="pl-c1">0.2</span>.<span class="pl-c1">9</span>
<span class="pl-s1">jaxlib</span>.<span class="pl-s1">__version__</span> <span class="pl-c1">=</span> <span class="pl-c1">0.1</span>.<span class="pl-c1">61</span></pre></div>