TypeError: Argument 'cpu:0' of type <class 'jaxlib.xla_extension.CpuDevice'> is not a valid JAX type.

[Joy-Preetha]

TypeError: Argument 'cpu:0' of type <class 'jaxlib.xla_extension.CpuDevice'> is not a valid JAX type.

I encountered this error while trying to run the below code.

def match_faces3(desc):
        distances = np.empty((len(desc), len(database)))
        f.write("Descriptors")
        f.write(str(desc))
        time1 = time.time()
        for i, descr in enumerate(desc):
            for j, identity in enumerate(database):
                dist = []
                for k, id_desc in enumerate(identity[1]):
                    dist.append(cosine_dist(descr, id_desc))
                distances[i][j] = dist[jnp.argmin(jnp.asarray(dist))]
        time2 = time.time() - time1
        print("time2",time2)
        distances=distances.tolist()
        return jnp.asarray(distances)

jax.device_put(descriptors,jax.devices()[0])
matches_jit = jit(match_faces3)
%timeit matches_jit(jax.devices()[0])

Desc is a multidimensional array

I tried everything but couldn't figure out the problem. Please help @mattjj

Originally posted by @Joy-Preetha in #4416 (comment)

[jakevdp]

When you call matches_jit(jax.devices()[0]), you are passing a jax device to a function that expects an array. You should try passing an array instead; for example:

desc = jax.device_put(descriptors,jax.devices()[0])
matches_jit = jit(match_faces3)
%timeit matches_jit(desc)

Also, as a side note – JAX's runtime model is asynchronous, so if you're timing operations you should use the block_until_ready() method to ensure the timings are accurate:

%timeit matches_jit(desc).block_until_ready()

See Asynchronous Dispatch for more information.

[Joy-Preetha]

When you call matches_jit(jax.devices()[0]), you are passing a jax device to a function that expects an array. You should try passing an array instead; for example:

desc = jax.device_put(descriptors,jax.devices()[0])
matches_jit = jit(match_faces3)
%timeit matches_jit(desc)

Also, as a side note – JAX's runtime model is asynchronous, so if you're timing operations you should use the block_until_ready() method to ensure the timings are accurate:

%timeit matches_jit(desc).block_until_ready()

See Asynchronous Dispatch for more information.

I tried this previously but got this Exception, so tried that way.

Exception: The numpy.ndarray conversion method array() was called on the JAX Tracer object Traced<ShapedArray(float32[512])>with<DynamicJaxprTrace(level=0/1)>.

This error can occur when a JAX Tracer object is passed to a raw numpy function, or a method on a numpy.ndarray object. You might want to check that you are using jnp together with import jax.numpy as jnp rather than using np via import numpy as np. If this error arises on a line that involves array indexing, like x[idx], it may be that the array being indexed x is a raw numpy.ndarray while the indices idx are a JAX Tracer instance; in that case, you can instead write jax.device_put(x)[idx].

[jakevdp]

It is hard to say for sure because the example you gave is incomplete, but I suspect cosine_dist() is not a JAX function, and so is not compatible with JIT.

[Joy-Preetha]

It is hard to say for sure because the example you gave is incomplete, but I suspect cosine_dist() is not a JAX function, and so is not compatible with JIT.

I changed cos into a jax function but the same error.

Actually this is the original code which I have run using numpy, Im using jax to reduce the time taken,

        distances = np.empty((len(descriptors), len(database)))
        f.write("Descriptors")
        f.write(str(descriptors))
        time1 = time.time()
        for i, desc in enumerate(descriptors):
            for j, identity in enumerate(database):
                dist = []
                for k, id_desc in enumerate(identity[1]):
                    dist.append(cosine_dist(desc, id_desc))
                distances[i][j] = dist[np.argmin(dist)]
        time2 = time.time() - time1
        print("time2",time2)
        f.write("Distances")
        f.write(str(distances))

[jakevdp]

Thanks for the further information. If you want further help debugging this, I'd suggest putting together a minimal reproducible example and including it here: that will take the guesswork out of helping find the cause of your error.

[Joy-Preetha]

[('8174',
[array([ 9.2, 9.69, 7.8, -2.6,
8.12, -5.39, 1.2, 2.2,
-4.2, -2.21, 5.78, 1.9,
-1.7, 1.05, -2.1, -1.6,
-1.0, -3.1, 3.3, -4.6],
dtype=float32)]),
('30819',
[array([ 1.85, -7.70, -2.0, 5.65,
-7.58, 3.3, -5.0, -7.25,
-1.31, 8.06, -9.40, -1.2,
1.4, -9.3, 7.8, 3.17,
-1.2, -5.6, 3.39, 1.1],
dtype=float32)]),
('41873',
[array([ 1.7, -4.2, -7.5, 5.5,
-5.6, -1.5, 5.2, -2.0,
-1.13, -1.96, -1.7, 5.0,
4.55, -8.0, 4.0, -9.86,
-1.41, -5.4, 4.7, -1.28],
dtype=float32)])]
This is how the database looks like, it is a pickle file.

descriptors = np.array([[-1.69,4.4,4.27,1.96,2.7,-5.73,-5.41,1.12,2.5,2.09,5.8,-8.7,6.7,4.1,7.4,6.3,9.7,2.4,6.4,3.3]])

This is how the descriptors look like

The cosine_func is a normal cos function.
Now I want to run the above code using jax.

[jakevdp]

Thanks - rather than describing your code in words (what does "normal cos function" mean?) it would be more helpful if you could include a complete code snippet, i.e. something that someone could copy and paste into a runtime and see the same results you are seeing. The minimal reproducible example link from my earlier comment is a helpful resource with more details on how to construct such an example.

[Joy-Preetha]

import numpy as np
import time,datetime
from scipy.spatial.distance import cosine
import jax
import jax.numpy as jnp

f = open("write_out.txt", 'w+')
def cosine_dist(x, y):
    return cosine(x, y) * 0.5

distances = np.empty((len(descriptors), len(database)))
        f.write("Descriptors")
        f.write(str(descriptors))
        time1 = time.time()
        for i, desc in enumerate(descriptors):
            for j, identity in enumerate(database):
                dist = []
                for k, id_desc in enumerate(identity[1]):
                    dist.append(cosine_dist(desc, id_desc))
                distances[i][j] = dist[np.argmin(dist)]
        time2 = time.time() - time1
        print("time2",time2)
        f.write("Distances")
        f.write(str(distances))

The jax which I tried to implement is

distances = np.empty((len(desc), len(database)))
f.write("Descriptors")
f.write(str(desc))
time1 = time.time()
for i, descr in enumerate(desc):
    for j, identity in enumerate(database):
        dist = []
        for k, id_desc in enumerate(identity[1]):
            dist.append(cosine_dist(descr, id_desc))
            distances[i][j] = dist[jnp.argmin(jnp.asarray(dist))]
time2 = time.time() - time1
print("time2",time2)
distances=distances.tolist()
f.write(str(distances))

When I just changed this Jax time was more, so I tried putting it in a jit function.

[jakevdp]

Thanks! This is much more clear. The issue is that scipy.spatial.distance.cosine is not a JAX function, and so is incompatible with JIT:

from scipy.spatial.distance import cosine
from jax import jit
import jax.numpy as jnp

jit(cosine)(jnp.arange(4), jnp.arange(4))

Exception: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[4])>with<DynamicJaxprTrace(level=0/1)>.

This error can occur when a JAX Tracer object is passed to a raw numpy function, or a method on a numpy.ndarray object. You might want to check that you are using `jnp` together with `import jax.numpy as jnp` rather than using `np` via `import numpy as np`. If this error arises on a line that involves array indexing, like `x[idx]`, it may be that the array being indexed `x` is a raw numpy.ndarray while the indices `idx` are a JAX Tracer instance; in that case, you can instead write `jax.device_put(x)[idx]`.

If you want to JIT-compile this computation, you'll have to implement the cosine distance using JAX functions; for example:

import numpy as np
import jax.numpy as jnp
from jax import jit
from scipy.spatial.distance import cosine

@jit
def jax_cosine(u, v):
  return 1 - jnp.dot(u, v) / (jnp.linalg.norm(u) * jnp.linalg.norm(v))

u = np.random.rand(100)
v = np.random.rand(100)

print(np.allclose(cosine(u, v), jax_cosine(u, v)))
# True

Additionally, if you want your implementation to have good performance, you should avoid iteration over array axes in favor of vectorized operations via broadcasting or through tools like vmap.

[Joy-Preetha]

Thank you
But now I encountered this error.

in match_faces3(desc)
12 for k, id_desc in enumerate(identity[1]):
13 dist.append(jax_cosine(descr, id_desc))
---> 14 distances[i][j] = dist[jnp.argmin(jnp.asarray(dist))]
15 time2 = time.time() - time1
16 print("time2",time2)

TypeError: list indices must be integers or slices, not DynamicJaxprTracer @jakevdp

[jakevdp]

That makes sense: you're trying to iterate over a traced array. As I mentioned, you should avoid explicit iteration in favor of either vectorized operations or vmap.

[jakevdp]

Actually, looking more closely the issue is that dist is a list, and you're indexing it with a traced quantity. But using a vectorized approach will remove the need to explicitly accumulate results in such a list, so that would still be my recommendation.

[Joy-Preetha]

Hi, I took your suggestion and tried to understand how vmap can be implemented, but couldnt figure it out. Can you please help me in it.