Valid DeviceArray slicing in an xmap analysis triggers obsolete slicing exception #6240
Closed
3 tasks done
Labels
bug
Something isn't working
Comments
|
Ok, I can get this to work and produce my expected result if I do the following: @jit
def phase_delay(lm, uvw, frequency):
out_dtype = np.result_type(lm, uvw, frequency, np.complex64)
one = lm.dtype.type(1.0)
constant = lm.dtype.type(minus_two_pi_over_c)
constant = constant * out_dtype.type(1j)
term = one - lm[0]**2 - lm[1]**2
term = np.where(term > 0.0, term, 0.0)
n = np.sqrt(term)
inner_prod = lm[0]*uvw[0] + lm[1]*uvw[1] + (n - one)*uvw[2]
return np.exp(constant*inner_prod*frequency)
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", ...],
["row", ...],
["chan"]),
out_axes=["source", "row", "chan"],
axis_resources={"row": "row"})
xphase(lm, uvw, freq)but in the phase_delay function, lm, uvw and frequency are now typed as an array of length 2, array of length 3 and a scalar respectively. So the primary input axes get removed when the array types are traced in the jitted function. Is this expected? |
|
After playing with vmap a bit more, I understand that the vectorised axis is not present when traced through the vectorised function. So I understand xmap to be behaving as vmap does. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
producing the following exception
This occurs on
jax==jax==0.2.11andjaxlib==0.1.64I don't think it is #620 as I'll explain at the end of the traceback:
See the full example below:
Please:
What I find interesting is that
arr_ndim == 0so 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?The text was updated successfully, but these errors were encountered: