Encountering AttributeError with jaxfluids and JAX 0.4.14

[jackyvava]

Problem

I'm using the jaxfluids library, which is a computational fluid dynamics library that requires JAX. I have installed JAX and jaxlib with versions jax=0.4.14 and jaxlib=0.4.14, respectively. However, when running the provided examples, I encountered the following AttributeError:

AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'

I noticed that this issue has been discussed in various contexts, including GPT-3's responses. Is there a recommended solution for resolving this issue with the mentioned versions of JAX?

Solution

After researching and experimenting, I found a workaround to resolve this issue. I used PyCharm's "Replace in file" feature to replace all instances of jnp.DeviceArray with jnp.Array. This helped me successfully run the examples without encountering the AttributeError.

Here's how I performed the replacement:

Open the file where jnp.DeviceArray is used.
Use the "Replace in file" feature in PyCharm.
Replace all occurrences of jnp.DeviceArray with jnp.Array.
Please note that this is a specific workaround I tried, and it may not be the only solution. However, it worked for me with the mentioned versions of JAX. Others who are facing similar issues might find this workaround helpful.

Result

Success

[weihuangsheng]

@jackyvava Genius solution! You are my God! I just faced the same problem and use the ”global search and replace“ function in VScode, which really help!

[trbedwards]

For me I had to replace with jax.Array and add the import jax statement if jax wasn't being imported

[dbezgin]

Sorry for the slow reply. We have updated jax.DeviceArray to jax.ndarray which should be compatible with newer JAX versions. Thanks for your input on this. I am closing this issue for now but please feel free to reopen.