<p dir="auto">Please:</p>
<ul dir="auto">
<li>[x ] Check for duplicate issues.</li>
<li>[ x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">This code generates a Traceback. If I replace <code class="notranslate">jax.opt.adam</code> with <code class="notranslate">jax_opt.sgd</code>, it seems to work fine. That is why I suspect something is wrong in <code class="notranslate">adam</code>.  This is in jax version 0.2.25.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.example_libraries.optimizers as jax_opt
from jax import value_and_grad
import jax.numpy as np

opt_init, opt_update, get_params = jax_opt.adam(1e-2)

opt_state = opt_init(np.array(1.1))

def objective(x):
    return (x**2).squeeze()

def step(i, opt_state):
    value, grads = value_and_grad(objective)(get_params(opt_state))
    return value, opt_update(step, grads, opt_state)

step(0, opt_state)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">example_libraries</span>.<span class="pl-s1">optimizers</span> <span class="pl-k">as</span> <span class="pl-s1">jax_opt</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">value_and_grad</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>

<span class="pl-s1">opt_init</span>, <span class="pl-s1">opt_update</span>, <span class="pl-s1">get_params</span> <span class="pl-c1">=</span> <span class="pl-s1">jax_opt</span>.<span class="pl-en">adam</span>(<span class="pl-c1">1e-2</span>)

<span class="pl-s1">opt_state</span> <span class="pl-c1">=</span> <span class="pl-en">opt_init</span>(<span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-c1">1.1</span>))

<span class="pl-k">def</span> <span class="pl-en">objective</span>(<span class="pl-s1">x</span>):
    <span class="pl-k">return</span> (<span class="pl-s1">x</span><span class="pl-c1">**</span><span class="pl-c1">2</span>).<span class="pl-en">squeeze</span>()

<span class="pl-k">def</span> <span class="pl-en">step</span>(<span class="pl-s1">i</span>, <span class="pl-s1">opt_state</span>):
    <span class="pl-s1">value</span>, <span class="pl-s1">grads</span> <span class="pl-c1">=</span> <span class="pl-en">value_and_grad</span>(<span class="pl-s1">objective</span>)(<span class="pl-en">get_params</span>(<span class="pl-s1">opt_state</span>))
    <span class="pl-k">return</span> <span class="pl-s1">value</span>, <span class="pl-en">opt_update</span>(<span class="pl-s1">step</span>, <span class="pl-s1">grads</span>, <span class="pl-s1">opt_state</span>)

<span class="pl-en">step</span>(<span class="pl-c1">0</span>, <span class="pl-s1">opt_state</span>)</pre></div>
<ul dir="auto">
<li>[x ] If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">I get this resulting error</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
&lt;ipython-input-7-caac38fbc2ef&gt; in &lt;module&gt;()
     14     return value, opt_update(step, grads, opt_state)
     15 
---&gt; 16 step(0, opt_state)

3 frames
/usr/local/lib/python3.7/dist-packages/jax/example_libraries/optimizers.py in update(i, g, state)
    408     m = (1 - b1) * g + b1 * m  # First  moment estimate.
    409     v = (1 - b2) * jnp.square(g) + b2 * v  # Second moment estimate.
--&gt; 410     mhat = m / (1 - jnp.asarray(b1, m.dtype) ** (i + 1))  # Bias correction.
    411     vhat = v / (1 - jnp.asarray(b2, m.dtype) ** (i + 1))
    412     x = x - step_size(i) * mhat / (jnp.sqrt(vhat) + eps)

TypeError: unsupported operand type(s) for +: 'function' and 'int'"><pre class="notranslate"><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span>
<span class="pl-v">TypeError</span>                                 <span class="pl-v">Traceback</span> (<span class="pl-s1">most</span> <span class="pl-s1">recent</span> <span class="pl-s1">call</span> <span class="pl-s1">last</span>)
<span class="pl-c1">&lt;</span><span class="pl-s1">ipython</span><span class="pl-c1">-</span><span class="pl-s1">input</span><span class="pl-c1">-</span><span class="pl-c1">7</span><span class="pl-c1">-</span><span class="pl-s1">caac38fbc2ef</span><span class="pl-c1">&gt;</span> <span class="pl-s1">in</span> <span class="pl-c1">&lt;</span><span class="pl-s1">module</span><span class="pl-c1">&gt;</span>()
     <span class="pl-c1">14</span>     <span class="pl-k">return</span> <span class="pl-s1">value</span>, <span class="pl-en">opt_update</span>(<span class="pl-s1">step</span>, <span class="pl-s1">grads</span>, <span class="pl-s1">opt_state</span>)
     <span class="pl-c1">15</span> 
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">16</span> <span class="pl-en">step</span>(<span class="pl-c1">0</span>, <span class="pl-s1">opt_state</span>)

<span class="pl-c1">3</span> <span class="pl-s1">frames</span>
<span class="pl-c1">/</span><span class="pl-s1">usr</span><span class="pl-c1">/</span><span class="pl-s1">local</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">7</span><span class="pl-c1">/</span><span class="pl-s1">dist</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">example_libraries</span><span class="pl-c1">/</span><span class="pl-s1">optimizers</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-en">update</span>(<span class="pl-s1">i</span>, <span class="pl-s1">g</span>, <span class="pl-s1">state</span>)
    <span class="pl-c1">408</span>     <span class="pl-s1">m</span> <span class="pl-c1">=</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">b1</span>) <span class="pl-c1">*</span> <span class="pl-s1">g</span> <span class="pl-c1">+</span> <span class="pl-s1">b1</span> <span class="pl-c1">*</span> <span class="pl-s1">m</span>  <span class="pl-c"># First  moment estimate.</span>
    <span class="pl-c1">409</span>     <span class="pl-s1">v</span> <span class="pl-c1">=</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">b2</span>) <span class="pl-c1">*</span> <span class="pl-s1">jnp</span>.<span class="pl-en">square</span>(<span class="pl-s1">g</span>) <span class="pl-c1">+</span> <span class="pl-s1">b2</span> <span class="pl-c1">*</span> <span class="pl-s1">v</span>  <span class="pl-c"># Second moment estimate.</span>
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">410</span>     <span class="pl-s1">mhat</span> <span class="pl-c1">=</span> <span class="pl-s1">m</span> <span class="pl-c1">/</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">jnp</span>.<span class="pl-en">asarray</span>(<span class="pl-s1">b1</span>, <span class="pl-s1">m</span>.<span class="pl-s1">dtype</span>) <span class="pl-c1">**</span> (<span class="pl-s1">i</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>))  <span class="pl-c"># Bias correction.</span>
    <span class="pl-c1">411</span>     <span class="pl-s1">vhat</span> <span class="pl-c1">=</span> <span class="pl-s1">v</span> <span class="pl-c1">/</span> (<span class="pl-c1">1</span> <span class="pl-c1">-</span> <span class="pl-s1">jnp</span>.<span class="pl-en">asarray</span>(<span class="pl-s1">b2</span>, <span class="pl-s1">m</span>.<span class="pl-s1">dtype</span>) <span class="pl-c1">**</span> (<span class="pl-s1">i</span> <span class="pl-c1">+</span> <span class="pl-c1">1</span>))
    <span class="pl-c1">412</span>     <span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">x</span> <span class="pl-c1">-</span> <span class="pl-en">step_size</span>(<span class="pl-s1">i</span>) <span class="pl-c1">*</span> <span class="pl-s1">mhat</span> <span class="pl-c1">/</span> (<span class="pl-s1">jnp</span>.<span class="pl-en">sqrt</span>(<span class="pl-s1">vhat</span>) <span class="pl-c1">+</span> <span class="pl-s1">eps</span>)

<span class="pl-v">TypeError</span>: <span class="pl-en">unsupported</span> <span class="pl-s1">operand</span> <span class="pl-s1">type</span>(<span class="pl-s1">s</span>) <span class="pl-s1">for</span> <span class="pl-c1">+</span>: <span class="pl-s">'function'</span> <span class="pl-c1">and</span> <span class="pl-s">'int'</span></pre></div>
<p dir="auto">The documentation of jax.devices says that it returns all devices. This is not true in multi-backend setups, it returns only the devices from the default backend. This was confusing to me initially, and I also encountered this confusion in issue <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="604212252" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/2785" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/2785/hovercard" href="https://github.com/google/jax/pull/2785">#2785</a></p>
<p dir="auto">The simplest change would be to the documentation: if no <code class="notranslate">backend</code> is specified then only the devices on the default backend are returned (along with an explanation of what is the default backend).</p>
<p dir="auto">A better change may be though to say that it returns all devices, in the order of backend priority, such that <code class="notranslate">devices()[0]</code> is the same as now. This may break some code though.</p>