"AttributeError: Unrecognized config option: jax_array" in windtunnel3d.py example

[astefanz]

Line 33 jax.config.update('jax_array', True) in windtunnel3d.py throws an exception. Looks it should be removed.

Executing the windtunnel3d.py example in a Google Colab instance throws this exception:

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

AttributeError                            Traceback (most recent call last)

[/content/XLB/examples/CFD/windtunnel3d.py](https://localhost:8080/#) in <module>
     31 # disable JIt compilation
     32 
---> 33 jax.config.update('jax_array', True)
     34 
     35 class Car(KBCSim):

[/usr/local/lib/python3.10/dist-packages/jax/_src/config.py](https://localhost:8080/#) in update(self, name, val)
     85   def update(self, name, val):
     86     if name not in self._value_holders:
---> 87       raise AttributeError(f"Unrecognized config option: {name}")
     88     self._value_holders[name]._set(val)
     89 

AttributeError: Unrecognized config option: jax_array

This happens on an L4 instance. I install dependencies by executing, in Colab notebook cells:

!pip install pyvista numpy matplotlib Rtree trimesh jmp orbax-checkpoint termcolor

And then cloning the repository and executing the example:

!git clone https://github.com/Autodesk/XLB

%run XLB/examples/CFD/windtunnel3d.py

Looks like the operation in line 33 was removed after the array migration in JAX 0.4.1, per: https://jax.readthedocs.io/en/latest/jax_array_migration.html and this line should just be removed. Thanks!

[hsalehipour]

Thanks for spotting this issue! You are right. We have a major refactoring coming in very soon (see branch major-refactoring) in which we have added a second backened to XLB based on Warp. We will fix this issue in that version. Please stay tuned @astefanz