Issue 1:
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.

xmap doesn't seem to like the following (although jit is fine with it), triggering what looks like #620.

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

producing the following exception

IndexError: Too many indices for array: 2 non-None/Ellipsis indices for dim 0.

This occurs on jax==jax==0.2.11 and jaxlib==0.1.64

I don't think it is #620 as I'll explain at the end of the traceback:

See the full example below:

Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

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__ == "__main__":
jax.config.update("jax_enable_x64", 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()) ("row",)):
xphase = xmap(phase_delay,
in_axes=(["source", "lm"],
["row", "uvw"],
["chan"]),
out_axes=["source", "row", "chan"],
axis_resources={"row": "row"})

xphase(lm, uvw, freq)

If applicable, include full error messages/tracebacks.

$ 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 <module>
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.

What I find interesting is that arr_ndim == 0 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?

Issue 2:
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 :)

by the way, the grad function works for different input values (e.g. x=[0.]*3)

Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug

Reproducible example:

in my drive because of two images but it's a colab notebook :)

traceback: (the print statements in the beginning are there to check the value of arr)

are there nans in arr? False
max Traced<ConcreteArray(1.1942956)>with<JVPTrace(level=2/0)>
with primal = DeviceArray(1.1942956, dtype=float32)
tangent = Traced<ShapedArray(float32[]):JaxprTrace(level=1/0)>
min Traced<ConcreteArray(-0.097952805)>with<JVPTrace(level=2/0)>
with primal = DeviceArray(-0.09795281, dtype=float32)
tangent = Traced<ShapedArray(float32[]):JaxprTrace(level=1/0)>

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

UnfilteredStackTrace Traceback (most recent call last)

<ipython-input-40-72c254e38e5c> in <module>()
2 config.update("jax_debug_nans", True)
----> 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:
--> 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):
--> 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:
--> 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:
--> 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(
-> 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:
--> 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)
--> 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)
--> 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:
--> 166 ans = self.f(*args, **dict(self.params, **kwargs))
167 except:

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

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

<ipython-input-34-fc9812f7e0ed> in jax_rgb2lab(img)
16 print('min', arr.min()) # debugging purposes
---> 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):
--> 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`."""
--> 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)
--> 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)
--> 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):
--> 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)
--> 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):
--> 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)
--> 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)
--> 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:
--> 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())):
--> 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)

<ipython-input-40-72c254e38e5c> in <module>()
1 from jax.config import config
2 config.update("jax_debug_nans", True)
----> 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)

<ipython-input-36-5e6c764e260d> 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)'''
---> 12 return dE(transformation(raw1, x), target1)
13
14 jit_loss = jax.jit(loss)

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

<ipython-input-34-fc9812f7e0ed> in jax_rgb2lab(img)
15 print('max', arr.max()) # debugging purposes
16 print('min', arr.min()) # debugging purposes
---> 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):
--> 639 return lax.pow(x1, x2)
640
641 # Integer power => use binary exponentiation.

FloatingPointError: invalid value (nan) encountered in pow

If applicable, include full error messages/tracebacks.