All input arrays must have the same shape

[PoulomiPradhan]

pt = jnp.array([0.,y,dist]) this is throwing ValueError "All input arrays must have the same shape" y is of type class <jax.interpreters.ad.JVPTracer> and dist is of type float.
out = stack([asarray(elt,dtype=dtype)]) in lax_numpy.py (line number 1900) ) is causing this.
How can I overcome this error ?

[jakevdp]

JAX does not support ragged arrays. So for example, this will work, because all entries in the array are the same size:

import jax.numpy as jnp
y = 1.
dist = 2.
pt = jnp.array([0., y, dist])
print(pt)
# [0. 1. 2.]

This, however, will raise an error because you're trying to construct a 2D array with rows of different sizes:

y = jnp.arange(4)
dist = jnp.ones(4)
pt = jnp.array([0., y, dist])
# ValueError: All input arrays must have the same shape.

To fix this, you can do a couple things depending on what you were hoping the output would be. For example, if you wanted a row of zeros, in your output array, you could do this:

pt = jnp.array([jnp.zeros_like(y), y, dist])
print(pt)
# [[0. 0. 0. 0.]
#  [0. 1. 2. 3.]
#  [1. 1. 1. 1.]]

Hope that helps!

[PoulomiPradhan]

Yeah. I am trying to rewrite files in rayoptics packages so that they support jax. Didn't find any package for Ray tracing which supports Jax.

[jakevdp]

Ah, that makes sense! The key thing then is to avoid any operations that will attempt to convert jax arrays to numpy arrays. The sign that that is happening is a TracerArrayConversionError.

[PoulomiPradhan]

Ok. Thanks I understnad! So there must be some functions in the rayoptics package those are still using numpy instead of jax.numpy.

[jakevdp]

Yes, for example rt.dot(x), if rt is a numpy array, would attempt to convert its argument to a numpy array.

[PoulomiPradhan]

Yes. I understand that. Thank you