<p dir="auto">I'm pretty excited about the new xmap feature and so I took it for a spin.I know xmap is a bit unstable at the moment (or I may have done something incorrectly!) so I thought I'd raise the following issue.</p>
<p dir="auto">xmap doesn't seem to like the following (although jit is fine with it), triggering what looks like <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="434056089" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/620" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/620/hovercard" href="https://github.com/google/jax/issues/620">#620</a>.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="    l = lm[:, 0, None, None]  # noqa"><pre class="notranslate">    <span class="pl-s1">l</span> <span class="pl-c1">=</span> <span class="pl-s1">lm</span>[:, <span class="pl-c1">0</span>, <span class="pl-c1">None</span>, <span class="pl-c1">None</span>]  <span class="pl-c"># noqa</span></pre></div>
<p dir="auto">producing the following exception</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="IndexError: Too many indices for array: 2 non-None/Ellipsis indices for dim 0."><pre class="notranslate"><code class="notranslate">IndexError: Too many indices for array: 2 non-None/Ellipsis indices for dim 0.
</code></pre></div>
<p dir="auto">This occurs on <code class="notranslate">jax==jax==0.2.11</code> and <code class="notranslate">jaxlib==0.1.64</code></p>
<p dir="auto">I don't think it is <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="434056089" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/620" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/620/hovercard" href="https://github.com/google/jax/issues/620">#620</a> as I'll explain at the end of the traceback:</p>
<p dir="auto">See the  full example below:</p>
<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>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.config
from jax.experimental.maps import xmap, mesh
import jax.numpy as np
from jax import jit
import numpy as onp

minus_two_pi_over_c = -2*np.pi/3e8

@jit
def phase_delay(lm, uvw, frequency):
    out_dtype = np.result_type(lm, uvw, frequency, np.complex64)

    one = lm.dtype.type(1.0)
    neg_two_pi_over_c = lm.dtype.type(minus_two_pi_over_c)
    complex_one = out_dtype.type(1j)

    l = lm[:, 0, None, None]  # noqa
    m = lm[:, 1, None, None]

    u = uvw[None, :, 0, None]
    v = uvw[None, :, 1, None]
    w = uvw[None, :, 2, None]

    n = np.sqrt(one - l**2 - m**2) - one

    real_phase = (neg_two_pi_over_c *
                  (l * u + m * v + n * w) *
                  frequency[None, None, :])

    return np.exp(complex_one*real_phase)

if __name__ == &quot;__main__&quot;:
    jax.config.update(&quot;jax_enable_x64&quot;, True)
    src = 10
    row = 1000
    chan = 64
    key = jax.random.PRNGKey(42)

    lm = jax.random.normal(key, (src, 2)) - 0.5
    uvw = (jax.random.normal(key, (row, 3)) - 0.5)*10000
    freq = np.linspace(.856e9, 2*.856e9, chan)

    with mesh(onp.array(jax.devices()) (&quot;row&quot;,)):
        xphase = xmap(phase_delay,
                      in_axes=([&quot;source&quot;, &quot;lm&quot;],
                               [&quot;row&quot;, &quot;uvw&quot;],
                               [&quot;chan&quot;]),
                      out_axes=[&quot;source&quot;, &quot;row&quot;, &quot;chan&quot;],
                      axis_resources={&quot;row&quot;: &quot;row&quot;})

        xphase(lm, uvw, freq)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">config</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span>.<span class="pl-s1">maps</span> <span class="pl-k">import</span> <span class="pl-s1">xmap</span>, <span class="pl-s1">mesh</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-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">jit</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">onp</span>

<span class="pl-s1">minus_two_pi_over_c</span> <span class="pl-c1">=</span> <span class="pl-c1">-</span><span class="pl-c1">2</span><span class="pl-c1">*</span><span class="pl-s1">np</span>.<span class="pl-s1">pi</span><span class="pl-c1">/</span><span class="pl-c1">3e8</span>

<span class="pl-en">@<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">phase_delay</span>(<span class="pl-s1">lm</span>, <span class="pl-s1">uvw</span>, <span class="pl-s1">frequency</span>):
    <span class="pl-s1">out_dtype</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">result_type</span>(<span class="pl-s1">lm</span>, <span class="pl-s1">uvw</span>, <span class="pl-s1">frequency</span>, <span class="pl-s1">np</span>.<span class="pl-s1">complex64</span>)

    <span class="pl-s1">one</span> <span class="pl-c1">=</span> <span class="pl-s1">lm</span>.<span class="pl-s1">dtype</span>.<span class="pl-en">type</span>(<span class="pl-c1">1.0</span>)
    <span class="pl-s1">neg_two_pi_over_c</span> <span class="pl-c1">=</span> <span class="pl-s1">lm</span>.<span class="pl-s1">dtype</span>.<span class="pl-en">type</span>(<span class="pl-s1">minus_two_pi_over_c</span>)
    <span class="pl-s1">complex_one</span> <span class="pl-c1">=</span> <span class="pl-s1">out_dtype</span>.<span class="pl-en">type</span>(<span class="pl-c1">1j</span>)

    <span class="pl-s1">l</span> <span class="pl-c1">=</span> <span class="pl-s1">lm</span>[:, <span class="pl-c1">0</span>, <span class="pl-c1">None</span>, <span class="pl-c1">None</span>]  <span class="pl-c"># noqa</span>
    <span class="pl-s1">m</span> <span class="pl-c1">=</span> <span class="pl-s1">lm</span>[:, <span class="pl-c1">1</span>, <span class="pl-c1">None</span>, <span class="pl-c1">None</span>]

    <span class="pl-s1">u</span> <span class="pl-c1">=</span> <span class="pl-s1">uvw</span>[<span class="pl-c1">None</span>, :, <span class="pl-c1">0</span>, <span class="pl-c1">None</span>]
    <span class="pl-s1">v</span> <span class="pl-c1">=</span> <span class="pl-s1">uvw</span>[<span class="pl-c1">None</span>, :, <span class="pl-c1">1</span>, <span class="pl-c1">None</span>]
    <span class="pl-s1">w</span> <span class="pl-c1">=</span> <span class="pl-s1">uvw</span>[<span class="pl-c1">None</span>, :, <span class="pl-c1">2</span>, <span class="pl-c1">None</span>]

    <span class="pl-s1">n</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">sqrt</span>(<span class="pl-s1">one</span> <span class="pl-c1">-</span> <span class="pl-s1">l</span><span class="pl-c1">**</span><span class="pl-c1">2</span> <span class="pl-c1">-</span> <span class="pl-s1">m</span><span class="pl-c1">**</span><span class="pl-c1">2</span>) <span class="pl-c1">-</span> <span class="pl-s1">one</span>

    <span class="pl-s1">real_phase</span> <span class="pl-c1">=</span> (<span class="pl-s1">neg_two_pi_over_c</span> <span class="pl-c1">*</span>
                  (<span class="pl-s1">l</span> <span class="pl-c1">*</span> <span class="pl-s1">u</span> <span class="pl-c1">+</span> <span class="pl-s1">m</span> <span class="pl-c1">*</span> <span class="pl-s1">v</span> <span class="pl-c1">+</span> <span class="pl-s1">n</span> <span class="pl-c1">*</span> <span class="pl-s1">w</span>) <span class="pl-c1">*</span>
                  <span class="pl-s1">frequency</span>[<span class="pl-c1">None</span>, <span class="pl-c1">None</span>, :])

    <span class="pl-k">return</span> <span class="pl-s1">np</span>.<span class="pl-en">exp</span>(<span class="pl-s1">complex_one</span><span class="pl-c1">*</span><span class="pl-s1">real_phase</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">jax</span>.<span class="pl-s1">config</span>.<span class="pl-en">update</span>(<span class="pl-s">"jax_enable_x64"</span>, <span class="pl-c1">True</span>)
    <span class="pl-s1">src</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>
    <span class="pl-s1">row</span> <span class="pl-c1">=</span> <span class="pl-c1">1000</span>
    <span class="pl-s1">chan</span> <span class="pl-c1">=</span> <span class="pl-c1">64</span>
    <span class="pl-s1">key</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">42</span>)

    <span class="pl-s1">lm</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">key</span>, (<span class="pl-s1">src</span>, <span class="pl-c1">2</span>)) <span class="pl-c1">-</span> <span class="pl-c1">0.5</span>
    <span class="pl-s1">uvw</span> <span class="pl-c1">=</span> (<span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">key</span>, (<span class="pl-s1">row</span>, <span class="pl-c1">3</span>)) <span class="pl-c1">-</span> <span class="pl-c1">0.5</span>)<span class="pl-c1">*</span><span class="pl-c1">10000</span>
    <span class="pl-s1">freq</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">linspace</span>(<span class="pl-c1">.856e9</span>, <span class="pl-c1">2</span><span class="pl-c1">*</span><span class="pl-c1">.856e9</span>, <span class="pl-s1">chan</span>)

    <span class="pl-k">with</span> <span class="pl-en">mesh</span>(<span class="pl-s1">onp</span>.<span class="pl-en">array</span>(<span class="pl-s1">jax</span>.<span class="pl-en">devices</span>()) (<span class="pl-s">"row"</span>,)):
        <span class="pl-s1">xphase</span> <span class="pl-c1">=</span> <span class="pl-en">xmap</span>(<span class="pl-s1">phase_delay</span>,
                      <span class="pl-s1">in_axes</span><span class="pl-c1">=</span>([<span class="pl-s">"source"</span>, <span class="pl-s">"lm"</span>],
                               [<span class="pl-s">"row"</span>, <span class="pl-s">"uvw"</span>],
                               [<span class="pl-s">"chan"</span>]),
                      <span class="pl-s1">out_axes</span><span class="pl-c1">=</span>[<span class="pl-s">"source"</span>, <span class="pl-s">"row"</span>, <span class="pl-s">"chan"</span>],
                      <span class="pl-s1">axis_resources</span><span class="pl-c1">=</span>{<span class="pl-s">"row"</span>: <span class="pl-s">"row"</span>})

        <span class="pl-en">xphase</span>(<span class="pl-s1">lm</span>, <span class="pl-s1">uvw</span>, <span class="pl-s1">freq</span>)</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="$ XLA_FLAGS=&quot;--xla_force_host_platform_device_count=8&quot; python  test_xmap.py 
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py:411: UserWarning: xmap is an experimental feature and probably has bugs!
  warn(&quot;xmap is an experimental feature and probably has bugs!&quot;)
Traceback (most recent call last):
  File &quot;test_xmap.py&quot;, line 95, in &lt;module&gt;
    xphase(lm, uvw, freq)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py&quot;, line 524, in fun_mapped
    backend=backend)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py&quot;, line 651, in bind
    return core.call_bind(self, fun, *args, **params)  # type: ignore
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py&quot;, line 1393, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py&quot;, line 654, in process
    return trace.process_xmap(self, fun, tracers, params)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py&quot;, line 600, in process_call
    return primitive.impl(f, *tracers, **params)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py&quot;, line 539, in xmap_impl
    axis_resources, resource_env, backend, *in_avals)(*args)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/linear_util.py&quot;, line 260, in memoized_fun
    ans = call(fun, *args)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py&quot;, line 554, in make_xmap_callable
    jaxpr, _, consts = pe.trace_to_jaxpr_final(fun, mapped_in_avals)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py&quot;, line 1228, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py&quot;, line 1208, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/linear_util.py&quot;, line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/api.py&quot;, line 382, in f_jitted
    return cpp_jitted_f(context, *args, **kwargs)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/api.py&quot;, line 351, in cache_miss_wrapper
    def cache_miss_wrapper(_, *args, **kw): return cache_miss(*args, **kw)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/api.py&quot;, line 301, in cache_miss
    donated_invars=donated_invars)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py&quot;, line 1402, in bind
    return call_bind(self, fun, *args, **params)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py&quot;, line 1393, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py&quot;, line 1405, in process
    return trace.process_call(self, fun, tracers, params)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py&quot;, line 1082, in process_call
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(f, self.main, in_avals)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py&quot;, line 1208, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/linear_util.py&quot;, line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File &quot;test_xmap.py&quot;, line 20, in phase_delay
    l = lm[:, 0, None, None]  # noqa
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py&quot;, line 527, in __getitem__
    def __getitem__(self, idx): return self.aval._getitem(self, idx)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 4462, in _rewriting_take
    return _gather(arr, treedef, static_idx, dynamic_idx)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 4469, in _gather
    indexer = _index_to_gather(shape(arr), idx)  # shared with _scatter_update
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 4553, in _index_to_gather
    idx = _canonicalize_tuple_index(len(x_shape), idx)
  File &quot;/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 4825, in _canonicalize_tuple_index
    raise IndexError(msg.format(len_without_none, arr_ndim))
IndexError: Too many indices for array: 2 non-None/Ellipsis indices for dim 0."><pre lang="traceback" class="notranslate"><code class="notranslate">$ XLA_FLAGS="--xla_force_host_platform_device_count=8" python  test_xmap.py 
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py:411: UserWarning: xmap is an experimental feature and probably has bugs!
  warn("xmap is an experimental feature and probably has bugs!")
Traceback (most recent call last):
  File "test_xmap.py", line 95, in &lt;module&gt;
    xphase(lm, uvw, freq)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py", line 524, in fun_mapped
    backend=backend)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py", line 651, in bind
    return core.call_bind(self, fun, *args, **params)  # type: ignore
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py", line 1393, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py", line 654, in process
    return trace.process_xmap(self, fun, tracers, params)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py", line 600, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py", line 539, in xmap_impl
    axis_resources, resource_env, backend, *in_avals)(*args)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/linear_util.py", line 260, in memoized_fun
    ans = call(fun, *args)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/experimental/maps.py", line 554, in make_xmap_callable
    jaxpr, _, consts = pe.trace_to_jaxpr_final(fun, mapped_in_avals)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 1228, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 1208, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/api.py", line 382, in f_jitted
    return cpp_jitted_f(context, *args, **kwargs)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/api.py", line 351, in cache_miss_wrapper
    def cache_miss_wrapper(_, *args, **kw): return cache_miss(*args, **kw)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/api.py", line 301, in cache_miss
    donated_invars=donated_invars)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py", line 1402, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py", line 1393, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py", line 1405, in process
    return trace.process_call(self, fun, tracers, params)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 1082, in process_call
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(f, self.main, in_avals)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 1208, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "test_xmap.py", line 20, in phase_delay
    l = lm[:, 0, None, None]  # noqa
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/core.py", line 527, in __getitem__
    def __getitem__(self, idx): return self.aval._getitem(self, idx)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py", line 4462, in _rewriting_take
    return _gather(arr, treedef, static_idx, dynamic_idx)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py", line 4469, in _gather
    indexer = _index_to_gather(shape(arr), idx)  # shared with _scatter_update
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py", line 4553, in _index_to_gather
    idx = _canonicalize_tuple_index(len(x_shape), idx)
  File "/home/sperkins/venv/jax/lib/python3.6/site-packages/jax/_src/numpy/lax_numpy.py", line 4825, in _canonicalize_tuple_index
    raise IndexError(msg.format(len_without_none, arr_ndim))
IndexError: Too many indices for array: 2 non-None/Ellipsis indices for dim 0.
</code></pre></div>
<p dir="auto">What I find interesting is that <code class="notranslate">arr_ndim == 0</code> so it looks to me like the tracing process is trying to slice a 0-length array. Thus I may be using xmap incorrectly, but I've done some checking. Any thoughts?</p>
<p dir="auto">Hi, I'm working in a simple optimization problem with image transformations. However, the gradient function fails and I can not understand why, as there is an array without nan values to the power of a number, which is not supposed to fail but gives nan. Furthermore, the numbers are not big. I would appreciate some help :)</p>
<p dir="auto">by the way, the grad function works for different input values (e.g. <code class="notranslate">x=[0.]*3</code>)</p>
<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</li>
</ul>
<p dir="auto">Reproducible example:<br>
<a href="https://drive.google.com/drive/folders/1rxjUF7nsagSyOKt4iGbM38d0CN5NgVne?usp=sharing" rel="nofollow">in my drive because of two images</a> but it's a colab notebook :)</p>
<p dir="auto">traceback: (the print statements in the beginning are there to check the value of <code class="notranslate">arr</code>)</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="are there nans in arr? False
max Traced&lt;ConcreteArray(1.1942956)&gt;with&lt;JVPTrace(level=2/0)&gt;
  with primal = DeviceArray(1.1942956, dtype=float32)
       tangent = Traced&lt;ShapedArray(float32[]):JaxprTrace(level=1/0)&gt;
min Traced&lt;ConcreteArray(-0.097952805)&gt;with&lt;JVPTrace(level=2/0)&gt;
  with primal = DeviceArray(-0.09795281, dtype=float32)
       tangent = Traced&lt;ShapedArray(float32[]):JaxprTrace(level=1/0)&gt;

---------------------------------------------------------------------------

UnfilteredStackTrace                      Traceback (most recent call last)

&lt;ipython-input-40-72c254e38e5c&gt; in &lt;module&gt;()
      2 config.update(&quot;jax_debug_nans&quot;, True)
----&gt; 3 dloss(raw1, target1, full_transformation, [.1]*3)
      4 

27 frames

/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    182     try:
--&gt; 183       return fun(*args, **kwargs)
    184     except Exception as e:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in grad_f(*args, **kwargs)
    828   def grad_f(*args, **kwargs):
--&gt; 829     _, g = value_and_grad_f(*args, **kwargs)
    830     return g

/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    182     try:
--&gt; 183       return fun(*args, **kwargs)
    184     except Exception as e:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in value_and_grad_f(*args, **kwargs)
    900     if not has_aux:
--&gt; 901       ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
    902     else:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in _vjp(fun, has_aux, reduce_axes, *primals)
   1996     out_primal, out_vjp = ad.vjp(
-&gt; 1997         flat_fun, primals_flat, reduce_axes=reduce_axes)
   1998     out_tree = out_tree()

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux, reduce_axes)
    114   if not has_aux:
--&gt; 115     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    116   else:

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
    101   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
--&gt; 102   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
    103   out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate)
    504     fun = trace_to_subjaxpr(fun, main, instantiate)
--&gt; 505     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    506     assert not env

/usr/local/lib/python3.7/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    165     try:
--&gt; 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:

&lt;ipython-input-36-5e6c764e260d&gt; in loss(raw1, target1, transformation, x)
     11   '''x represents the parameters of the transformation (powers in our case)'''
---&gt; 12   return dE(transformation(raw1, x), target1)
     13 

&lt;ipython-input-34-fc9812f7e0ed&gt; in dE(img1, img2)
     31 def dE(img1, img2):
---&gt; 32   return jnp.mean(jnp.sqrt(jnp.sum((jax_rgb2lab(img1) - jax_rgb2lab(img2))**2, axis=-1)))
     33 

&lt;ipython-input-34-fc9812f7e0ed&gt; in jax_rgb2lab(img)
     16     print('min', arr.min())  # debugging purposes
---&gt; 17     arr = jnp.power((arr + 0.055) / 1.055, 2.4) * mask + (arr / 12.92) * (1 - mask)
     18     arr = arr @ xyz_from_rgb_T

/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py in power(x1, x2)
    638   if not issubdtype(dtype, integer):
--&gt; 639     return lax.pow(x1, x2)
    640 

/usr/local/lib/python3.7/dist-packages/jax/_src/lax/lax.py in pow(x, y)
    299   r&quot;&quot;&quot;Elementwise power: :math:`x^y`.&quot;&quot;&quot;
--&gt; 300   return pow_p.bind(x, y)
    301 

/usr/local/lib/python3.7/dist-packages/jax/core.py in bind(self, *args, **params)
    263     tracers = map(top_trace.full_raise, args)
--&gt; 264     out = top_trace.process_primitive(self, tracers, params)
    265     return map(full_lower, out) if self.multiple_results else full_lower(out)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
    282       raise NotImplementedError(msg)
--&gt; 283     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    284     if primitive.multiple_results:

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in standard_jvp2(jvprules, primitive, primals, tangents, **params)
    470 def standard_jvp2(jvprules, primitive, primals, tangents, **params):
--&gt; 471   val_out = primitive.bind(*primals, **params)
    472   tangents_out = (rule(t, val_out, *primals, **params) for rule, t in zip(jvprules, tangents)

/usr/local/lib/python3.7/dist-packages/jax/core.py in bind(self, *args, **params)
    263     tracers = map(top_trace.full_raise, args)
--&gt; 264     out = top_trace.process_primitive(self, tracers, params)
    265     return map(full_lower, out) if self.multiple_results else full_lower(out)

/usr/local/lib/python3.7/dist-packages/jax/core.py in process_primitive(self, primitive, tracers, params)
    602   def process_primitive(self, primitive, tracers, params):
--&gt; 603     return primitive.impl(*tracers, **params)
    604 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    248   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
--&gt; 249   return compiled_fun(*args)
    250 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in _execute_compiled_primitive(prim, compiled, result_handler, *args)
    365   out_bufs = compiled.execute(input_bufs)
--&gt; 366   check_special(prim.name, out_bufs)
    367   return result_handler(*out_bufs)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in check_special(name, bufs)
    384     for buf in bufs:
--&gt; 385       _check_special(name, buf.xla_shape(), buf)
    386 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
    390     if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--&gt; 391       raise FloatingPointError(f&quot;invalid value (nan) encountered in {name}&quot;)
    392     if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):

UnfilteredStackTrace: FloatingPointError: invalid value (nan) encountered in pow

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------


The above exception was the direct cause of the following exception:

FloatingPointError                        Traceback (most recent call last)

&lt;ipython-input-40-72c254e38e5c&gt; in &lt;module&gt;()
      1 from jax.config import config
      2 config.update(&quot;jax_debug_nans&quot;, True)
----&gt; 3 dloss(raw1, target1, full_transformation, [.1]*3)
      4 
      5 # jax.make_jaxpr(dloss, static_argnums=[0, 1, 2])(raw1, target1, full_transformation, [.1]*3)

&lt;ipython-input-36-5e6c764e260d&gt; in loss(raw1, target1, transformation, x)
     10 def loss(raw1, target1, transformation, x):
     11   '''x represents the parameters of the transformation (powers in our case)'''
---&gt; 12   return dE(transformation(raw1, x), target1)
     13 
     14 jit_loss = jax.jit(loss)

&lt;ipython-input-34-fc9812f7e0ed&gt; in dE(img1, img2)
     30 
     31 def dE(img1, img2):
---&gt; 32   return jnp.mean(jnp.sqrt(jnp.sum((jax_rgb2lab(img1) - jax_rgb2lab(img2))**2, axis=-1)))
     33 

&lt;ipython-input-34-fc9812f7e0ed&gt; in jax_rgb2lab(img)
     15     print('max', arr.max())  # debugging purposes
     16     print('min', arr.min())  # debugging purposes
---&gt; 17     arr = jnp.power((arr + 0.055) / 1.055, 2.4) * mask + (arr / 12.92) * (1 - mask)
     18     arr = arr @ xyz_from_rgb_T
     19     # scale by CIE XYZ tristimulus values of the reference white point

/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py in power(x1, x2)
    637   dtype = _dtype(x1)
    638   if not issubdtype(dtype, integer):
--&gt; 639     return lax.pow(x1, x2)
    640 
    641   # Integer power =&gt; use binary exponentiation.

FloatingPointError: invalid value (nan) encountered in pow"><pre class="notranslate"><code class="notranslate">are there nans in arr? False
max Traced&lt;ConcreteArray(1.1942956)&gt;with&lt;JVPTrace(level=2/0)&gt;
  with primal = DeviceArray(1.1942956, dtype=float32)
       tangent = Traced&lt;ShapedArray(float32[]):JaxprTrace(level=1/0)&gt;
min Traced&lt;ConcreteArray(-0.097952805)&gt;with&lt;JVPTrace(level=2/0)&gt;
  with primal = DeviceArray(-0.09795281, dtype=float32)
       tangent = Traced&lt;ShapedArray(float32[]):JaxprTrace(level=1/0)&gt;

---------------------------------------------------------------------------

UnfilteredStackTrace                      Traceback (most recent call last)

&lt;ipython-input-40-72c254e38e5c&gt; in &lt;module&gt;()
      2 config.update("jax_debug_nans", True)
----&gt; 3 dloss(raw1, target1, full_transformation, [.1]*3)
      4 

27 frames

/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    182     try:
--&gt; 183       return fun(*args, **kwargs)
    184     except Exception as e:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in grad_f(*args, **kwargs)
    828   def grad_f(*args, **kwargs):
--&gt; 829     _, g = value_and_grad_f(*args, **kwargs)
    830     return g

/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    182     try:
--&gt; 183       return fun(*args, **kwargs)
    184     except Exception as e:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in value_and_grad_f(*args, **kwargs)
    900     if not has_aux:
--&gt; 901       ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
    902     else:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in _vjp(fun, has_aux, reduce_axes, *primals)
   1996     out_primal, out_vjp = ad.vjp(
-&gt; 1997         flat_fun, primals_flat, reduce_axes=reduce_axes)
   1998     out_tree = out_tree()

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux, reduce_axes)
    114   if not has_aux:
--&gt; 115     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    116   else:

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
    101   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
--&gt; 102   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
    103   out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate)
    504     fun = trace_to_subjaxpr(fun, main, instantiate)
--&gt; 505     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    506     assert not env

/usr/local/lib/python3.7/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    165     try:
--&gt; 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:

&lt;ipython-input-36-5e6c764e260d&gt; in loss(raw1, target1, transformation, x)
     11   '''x represents the parameters of the transformation (powers in our case)'''
---&gt; 12   return dE(transformation(raw1, x), target1)
     13 

&lt;ipython-input-34-fc9812f7e0ed&gt; in dE(img1, img2)
     31 def dE(img1, img2):
---&gt; 32   return jnp.mean(jnp.sqrt(jnp.sum((jax_rgb2lab(img1) - jax_rgb2lab(img2))**2, axis=-1)))
     33 

&lt;ipython-input-34-fc9812f7e0ed&gt; in jax_rgb2lab(img)
     16     print('min', arr.min())  # debugging purposes
---&gt; 17     arr = jnp.power((arr + 0.055) / 1.055, 2.4) * mask + (arr / 12.92) * (1 - mask)
     18     arr = arr @ xyz_from_rgb_T

/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py in power(x1, x2)
    638   if not issubdtype(dtype, integer):
--&gt; 639     return lax.pow(x1, x2)
    640 

/usr/local/lib/python3.7/dist-packages/jax/_src/lax/lax.py in pow(x, y)
    299   r"""Elementwise power: :math:`x^y`."""
--&gt; 300   return pow_p.bind(x, y)
    301 

/usr/local/lib/python3.7/dist-packages/jax/core.py in bind(self, *args, **params)
    263     tracers = map(top_trace.full_raise, args)
--&gt; 264     out = top_trace.process_primitive(self, tracers, params)
    265     return map(full_lower, out) if self.multiple_results else full_lower(out)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
    282       raise NotImplementedError(msg)
--&gt; 283     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    284     if primitive.multiple_results:

/usr/local/lib/python3.7/dist-packages/jax/interpreters/ad.py in standard_jvp2(jvprules, primitive, primals, tangents, **params)
    470 def standard_jvp2(jvprules, primitive, primals, tangents, **params):
--&gt; 471   val_out = primitive.bind(*primals, **params)
    472   tangents_out = (rule(t, val_out, *primals, **params) for rule, t in zip(jvprules, tangents)

/usr/local/lib/python3.7/dist-packages/jax/core.py in bind(self, *args, **params)
    263     tracers = map(top_trace.full_raise, args)
--&gt; 264     out = top_trace.process_primitive(self, tracers, params)
    265     return map(full_lower, out) if self.multiple_results else full_lower(out)

/usr/local/lib/python3.7/dist-packages/jax/core.py in process_primitive(self, primitive, tracers, params)
    602   def process_primitive(self, primitive, tracers, params):
--&gt; 603     return primitive.impl(*tracers, **params)
    604 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    248   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
--&gt; 249   return compiled_fun(*args)
    250 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in _execute_compiled_primitive(prim, compiled, result_handler, *args)
    365   out_bufs = compiled.execute(input_bufs)
--&gt; 366   check_special(prim.name, out_bufs)
    367   return result_handler(*out_bufs)

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in check_special(name, bufs)
    384     for buf in bufs:
--&gt; 385       _check_special(name, buf.xla_shape(), buf)
    386 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in _check_special(name, xla_shape, buf)
    390     if config.jax_debug_nans and np.any(np.isnan(buf.to_py())):
--&gt; 391       raise FloatingPointError(f"invalid value (nan) encountered in {name}")
    392     if config.jax_debug_infs and np.any(np.isinf(buf.to_py())):

UnfilteredStackTrace: FloatingPointError: invalid value (nan) encountered in pow

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------


The above exception was the direct cause of the following exception:

FloatingPointError                        Traceback (most recent call last)

&lt;ipython-input-40-72c254e38e5c&gt; in &lt;module&gt;()
      1 from jax.config import config
      2 config.update("jax_debug_nans", True)
----&gt; 3 dloss(raw1, target1, full_transformation, [.1]*3)
      4 
      5 # jax.make_jaxpr(dloss, static_argnums=[0, 1, 2])(raw1, target1, full_transformation, [.1]*3)

&lt;ipython-input-36-5e6c764e260d&gt; in loss(raw1, target1, transformation, x)
     10 def loss(raw1, target1, transformation, x):
     11   '''x represents the parameters of the transformation (powers in our case)'''
---&gt; 12   return dE(transformation(raw1, x), target1)
     13 
     14 jit_loss = jax.jit(loss)

&lt;ipython-input-34-fc9812f7e0ed&gt; in dE(img1, img2)
     30 
     31 def dE(img1, img2):
---&gt; 32   return jnp.mean(jnp.sqrt(jnp.sum((jax_rgb2lab(img1) - jax_rgb2lab(img2))**2, axis=-1)))
     33 

&lt;ipython-input-34-fc9812f7e0ed&gt; in jax_rgb2lab(img)
     15     print('max', arr.max())  # debugging purposes
     16     print('min', arr.min())  # debugging purposes
---&gt; 17     arr = jnp.power((arr + 0.055) / 1.055, 2.4) * mask + (arr / 12.92) * (1 - mask)
     18     arr = arr @ xyz_from_rgb_T
     19     # scale by CIE XYZ tristimulus values of the reference white point

/usr/local/lib/python3.7/dist-packages/jax/_src/numpy/lax_numpy.py in power(x1, x2)
    637   dtype = _dtype(x1)
    638   if not issubdtype(dtype, integer):
--&gt; 639     return lax.pow(x1, x2)
    640 
    641   # Integer power =&gt; use binary exponentiation.

FloatingPointError: invalid value (nan) encountered in pow
</code></pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> If applicable, include full error messages/tracebacks.</li>
</ul>