deepxde AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'

[Chaztikov]

Hi, I am using the anaconda version of deepxde (can provide additional details if necessary), I am getting this error

deepxde AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'

I see that conda list does not show deepxde among its packages. Is it inadvisable to use the conda version? I --force-reinstall'ed deepxde and the issue persists. Any suggestions are appreciated. Thanks!

Here is the full error message

import deepxde

File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/deepxde/init.py", line 17, in
from . import backend
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/deepxde/backend/init.py", line 100, in
load_backend(get_preferred_backend())
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/deepxde/backend/init.py", line 36, in load_backend
mod = importlib.import_module(".%s" % mod_name.replace(".", "_"), name)
File "/usr/local/anaconda/anaconda3/lib/python3.9/importlib/init.py", line 127, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/deepxde/backend/tensorflow/init.py", line 5, in
from .tensor import * # pylint: disable=redefined-builtin
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/deepxde/backend/tensorflow/tensor.py", line 4, in
import tensorflow as tf
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/init.py", line 51, in
from ._api.v2 import compat
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/_api/v2/compat/init.py", line 37, in
from . import v1
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/init.py", line 30, in
from . import compat
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/init.py", line 37, in
from . import v1
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/v1/init.py", line 47, in
from tensorflow._api.v2.compat.v1 import lite
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/init.py", line 9, in
from . import experimental
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/init.py", line 8, in
from . import authoring
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/authoring/init.py", line 8, in
from tensorflow.lite.python.authoring.authoring import compatible
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/lite/python/authoring/authoring.py", line 43, in
from tensorflow.lite.python import convert
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/lite/python/convert.py", line 28, in
from tensorflow.lite.python import util
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/tensorflow/lite/python/util.py", line 55, in
from jax import xla_computation as _xla_computation
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/jax/init.py", line 116, in
from .experimental.maps import soft_pmap as soft_pmap
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/jax/experimental/maps.py", line 26, in
from .. import numpy as jnp
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/jax/numpy/init.py", line 19, in
from . import fft as fft
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/jax/numpy/fft.py", line 17, in
from jax._src.numpy.fft import (
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/jax/_src/numpy/fft.py", line 19, in
from jax import lax
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/jax/lax/init.py", line 332, in
from jax._src.lax.fft import (
File "/afs/cas.unc.edu/users/c/t/ctalbot/.local/lib/python3.9/site-packages/jax/_src/lax/fft.py", line 145, in
xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft
AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'

[lululxvi]

The issue seems from jaxlib. Do you use JAX backend?

[Chaztikov]

Hi Professor Lu, Sorry for my late response, I'm honored to have yours.

I do not use the JAX backend (though I look forward to doing so)

It appears to be related to an issue related to my own recklessness: deepxde was installed locally on the cluster (for my username) but this was not identified by the cluster's install of anaconda3.

I'd say that this issue is resolved.

Thanks!