Module not found after freshly installation

[Qazalbash]

Description

After running the command,

pip install jaxampler

I am getting,

gradf@beepboop:~$ pip3 install jaxampler
Defaulting to user installation because normal site-packages is not writeable
Requirement already satisfied: jaxampler in /usr/local/lib/python3.10/dist-packages (0.0.6)
Requirement already satisfied: jax>=0.4.0 in ./.local/lib/python3.10/site-packages (from jaxampler) (0.4.23)
Requirement already satisfied: jaxlib>=0.4.0 in ./.local/lib/python3.10/site-packages (from jaxampler) (0.4.23)
Requirement already satisfied: typing-extensions>=4.5.0 in ./.local/lib/python3.10/site-packages (from jaxampler) (4.9.0)
Requirement already satisfied: jaxtyping>=0.2.24 in ./.local/lib/python3.10/site-packages (from jaxampler) (0.2.25)
Requirement already satisfied: matplotlib>=3.8.0 in ./.local/lib/python3.10/site-packages (from jaxampler) (3.8.2)
Requirement already satisfied: setuptools in ./.local/lib/python3.10/site-packages (from jaxampler) (69.0.3)
Requirement already satisfied: tfp-nightly in ./.local/lib/python3.10/site-packages (from jaxampler) (0.24.0.dev20240127)
Requirement already satisfied: tqdm in ./.local/lib/python3.10/site-packages (from jaxampler) (4.66.1)
Requirement already satisfied: twine in ./.local/lib/python3.10/site-packages (from jaxampler) (4.0.2)
Requirement already satisfied: wheel in ./.local/lib/python3.10/site-packages (from jaxampler) (0.42.0)
Requirement already satisfied: ml-dtypes>=0.2.0 in ./.local/lib/python3.10/site-packages (from jax>=0.4.0->jaxampler) (0.3.2)
Requirement already satisfied: numpy>=1.22 in ./.local/lib/python3.10/site-packages (from jax>=0.4.0->jaxampler) (1.26.3)
Requirement already satisfied: opt-einsum in ./.local/lib/python3.10/site-packages (from jax>=0.4.0->jaxampler) (3.3.0)
Requirement already satisfied: scipy>=1.9 in ./.local/lib/python3.10/site-packages (from jax>=0.4.0->jaxampler) (1.12.0)
Requirement already satisfied: typeguard<3,>=2.13.3 in ./.local/lib/python3.10/site-packages (from jaxtyping>=0.2.24->jaxampler) (2.13.3)
Requirement already satisfied: contourpy>=1.0.1 in ./.local/lib/python3.10/site-packages (from matplotlib>=3.8.0->jaxampler) (1.2.0)
Requirement already satisfied: cycler>=0.10 in ./.local/lib/python3.10/site-packages (from matplotlib>=3.8.0->jaxampler) (0.12.1)
Requirement already satisfied: fonttools>=4.22.0 in ./.local/lib/python3.10/site-packages (from matplotlib>=3.8.0->jaxampler) (4.47.2)
Requirement already satisfied: kiwisolver>=1.3.1 in ./.local/lib/python3.10/site-packages (from matplotlib>=3.8.0->jaxampler) (1.4.5)
Requirement already satisfied: packaging>=20.0 in /usr/lib/python3/dist-packages (from matplotlib>=3.8.0->jaxampler) (21.3)
Requirement already satisfied: pillow>=8 in /usr/lib/python3/dist-packages (from matplotlib>=3.8.0->jaxampler) (9.0.1)
Requirement already satisfied: pyparsing>=2.3.1 in /usr/lib/python3/dist-packages (from matplotlib>=3.8.0->jaxampler) (2.4.7)
Requirement already satisfied: python-dateutil>=2.7 in ./.local/lib/python3.10/site-packages (from matplotlib>=3.8.0->jaxampler) (2.8.2)
Requirement already satisfied: absl-py in ./.local/lib/python3.10/site-packages (from tfp-nightly->jaxampler) (2.1.0)
Requirement already satisfied: six>=1.10.0 in /usr/lib/python3/dist-packages (from tfp-nightly->jaxampler) (1.16.0)
Requirement already satisfied: decorator in ./.local/lib/python3.10/site-packages (from tfp-nightly->jaxampler) (5.1.1)
Requirement already satisfied: cloudpickle>=1.3 in ./.local/lib/python3.10/site-packages (from tfp-nightly->jaxampler) (3.0.0)
Requirement already satisfied: gast>=0.3.2 in ./.local/lib/python3.10/site-packages (from tfp-nightly->jaxampler) (0.5.4)
Requirement already satisfied: dm-tree in ./.local/lib/python3.10/site-packages (from tfp-nightly->jaxampler) (0.1.8)
Requirement already satisfied: pkginfo>=1.8.1 in ./.local/lib/python3.10/site-packages (from twine->jaxampler) (1.9.6)
Requirement already satisfied: readme-renderer>=35.0 in ./.local/lib/python3.10/site-packages (from twine->jaxampler) (42.0)
Requirement already satisfied: requests>=2.20 in /usr/lib/python3/dist-packages (from twine->jaxampler) (2.25.1)
Requirement already satisfied: requests-toolbelt!=0.9.0,>=0.8.0 in ./.local/lib/python3.10/site-packages (from twine->jaxampler) (1.0.0)
Requirement already satisfied: urllib3>=1.26.0 in /usr/lib/python3/dist-packages (from twine->jaxampler) (1.26.5)
Requirement already satisfied: importlib-metadata>=3.6 in /usr/lib/python3/dist-packages (from twine->jaxampler) (4.6.4)
Requirement already satisfied: keyring>=15.1 in /usr/lib/python3/dist-packages (from twine->jaxampler) (23.5.0)
Requirement already satisfied: rfc3986>=1.4.0 in ./.local/lib/python3.10/site-packages (from twine->jaxampler) (2.0.0)
Requirement already satisfied: rich>=12.0.0 in ./.local/lib/python3.10/site-packages (from twine->jaxampler) (13.7.0)
Requirement already satisfied: nh3>=0.2.14 in ./.local/lib/python3.10/site-packages (from readme-renderer>=35.0->twine->jaxampler) (0.2.15)
Requirement already satisfied: docutils>=0.13.1 in ./.local/lib/python3.10/site-packages (from readme-renderer>=35.0->twine->jaxampler) (0.20.1)
Requirement already satisfied: Pygments>=2.5.1 in ./.local/lib/python3.10/site-packages (from readme-renderer>=35.0->twine->jaxampler) (2.17.2)
Requirement already satisfied: markdown-it-py>=2.2.0 in ./.local/lib/python3.10/site-packages (from rich>=12.0.0->twine->jaxampler) (3.0.0)
Requirement already satisfied: mdurl~=0.1 in ./.local/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich>=12.0.0->twine->jaxampler) (0.1.2)
gradf@beepboop:~$ python3
Python 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jaxampler
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'jaxampler'
>>> 

site-packages contain these folder,

gradf@beepboop:~$ ls /home/gradf/.local/lib/python3.10/site-packages
absl                         exceptiongroup                  markdown_it                     pandas                     pytest-7.4.4.dist-info             tfp_nightly-0.24.0.dev20240127.dist-info
absl_py-2.1.0.dist-info      exceptiongroup-1.2.0.dist-info  markdown_it_py-3.0.0.dist-info  pandas-2.2.0.dist-info     python_dateutil-2.8.2.dist-info    tomli
cloudpickle                  fontTools                       matplotlib                      pip                        readme_renderer                    tomli-2.0.1.dist-info
cloudpickle-3.0.0.dist-info  fonttools-4.47.2.dist-info      matplotlib-3.8.2.dist-info      pip-23.3.2.dist-info       readme_renderer-42.0.dist-info     tqdm
contourpy                    gast                            mdurl                           pkginfo                    requests_toolbelt                  tqdm-4.66.1.dist-info
contourpy-1.2.0.dist-info    gast-0.5.4.dist-info            mdurl-0.1.2.dist-info           pkginfo-1.9.6.dist-info    requests_toolbelt-1.0.0.dist-info  tree
cycler                       iniconfig                       ml_dtypes                       pkg_resources              rfc3986                            twine
cycler-0.12.1.dist-info      iniconfig-2.0.0.dist-info       ml_dtypes-0.3.2.dist-info       pluggy                     rfc3986-2.0.0.dist-info            twine-4.0.2.dist-info
dateutil                     jax                             mpl_toolkits                    pluggy-1.4.0.dist-info     rich                               typeguard
decorator-5.1.1.dist-info    jax-0.4.23.dist-info            nh3                             __pycache__                rich-13.7.0.dist-info              typeguard-2.13.3.dist-info
decorator.py                 jaxlib                          nh3-0.2.15.dist-info            pygments                   scipy                              typing_extensions-4.9.0.dist-info
_distutils_hack              jaxlib-0.4.23.dist-info         numpy                           pygments-2.17.2.dist-info  scipy-1.12.0.dist-info             typing_extensions.py
distutils-precedence.pth     jaxtyping                       numpy-1.26.3.dist-info          pylab.py                   scipy.libs                         tzdata
dm_tree-0.1.8.dist-info      jaxtyping-0.2.25.dist-info      numpy.libs                      py.py                      setuptools                         tzdata-2023.4.dist-info
docutils                     kiwisolver                      opt_einsum                      _pytest                    setuptools-69.0.3.dist-info        wheel
docutils-0.20.1.dist-info    kiwisolver-1.4.5.dist-info      opt_einsum-3.3.0.dist-info      pytest                     tensorflow_probability             wheel-0.42.0.dist-info

This answer can be helpful in dealing with this issue.

What jaxampler version are you using?

v0.0.6

Which accelerator(s) are you using?

CPU

Additional system info?

1.26.3 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] uname_result(system='Linux', node='beepboop', release='5.15.0-76-generic', version='#83-Ubuntu SMP Thu Jun 15 19:16:32 UTC 2023', machine='x86_64')

NVIDIA GPU info

No response