search

Jammy2211 / PyAutoLens

PyAutoLens: Open Source Strong Gravitational Lensing
https://pyautolens.readthedocs.io/
MIT License
162 stars 32 forks source link

JAX on older HPC issues #258

Closed AstroAaron closed 6 months ago

AstroAaron commented 6 months ago

Hey, I get the following error when trying to run PyAutoLens on older nodes on an HPC that I am using. 

Traceback (most recent call last):
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/source_pixelized_2Lenses.py", line 63, in <module>
    import autofit as af
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/PyAutoFit/autofit/__init__.py", line 7, in <module>
    from .non_linear.grid.grid_search import GridSearch as SearchGridSearch
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/PyAutoFit/autofit/non_linear/grid/grid_search/__init__.py", line 9, in <module>
    from autofit.mapper import prior as p
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/PyAutoFit/autofit/mapper/__init__.py", line 1, in <module>
    from .prior.abstract import Prior
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/PyAutoFit/autofit/mapper/prior/__init__.py", line 1, in <module>
    from .gaussian import GaussianPrior
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/PyAutoFit/autofit/mapper/prior/gaussian.py", line 3, in <module>
    from autofit.jax_wrapper import register_pytree_node_class
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/PyAutoFit/autofit/jax_wrapper.py", line 32, in <module>
    from jax._src.tree_util import (
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/autolensNN/lib/python3.8/site-packages/jax/__init__.py", line 35, in <module>
    from jax import config as _config_module
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/autolensNN/lib/python3.8/site-packages/jax/config.py", line 17, in <module>
    from jax._src.config import config  # noqa: F401
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/autolensNN/lib/python3.8/site-packages/jax/_src/config.py", line 24, in <module>
    from jax._src import lib
  File "/projects/ag-riechers/users/aaron/ALMACE/PyAutoLensNN/autolensNN/lib/python3.8/site-packages/jax/_src/lib/__init__.py", line 84, in <module>
    cpu_feature_guard.check_cpu_features()
RuntimeError: This version of jaxlib was built using AVX instructions, which your CPU and/or operating system do not support. You may be able work around this issue by building jaxlib from source.

WIll the code still run properly if I downgrade JAX/build it from source as mentioned in https://stackoverflow.com/questions/77747723/runtimeerror-this-version-of-jaxlib-was-built-using-avx-instructions?

Or alternatively, is there a way of disabling JAX, as the nodes I am using do not have any GPUs anyway?

Jammy2211 commented 6 months ago

Ooops, not sure how you ended up with a JAX install but this was not intended. I think it may be related to a few files that I've now fixed.

Anyway, to dsiable this run:

pip uninstall jax pip uninstall jaxlib

JAX support is coming sooner or later... but its defo not ready yet ;)

AstroAaron commented 6 months ago

Thanks that worked!