pip install failed and got UnicodeDecodeError for jax>=0.3.15 on Windows, Python 3.9.13

[Jasonnor]

Description

I got UnicodeDecodeError when running pip install for jax 0.3.15 on Windows 10:

PS C:\tensorflow> pip install --no-cache-dir jax==0.3.15
Collecting jax==0.3.15
  Downloading jax-0.3.15.tar.gz (1.0 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.0/1.0 MB 3.5 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [6 lines of output]
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "C:\Users\Jason Wu\AppData\Local\Temp\pip-install-5rpko2f6\jax_4375f8ec13e046228768476f1ee3a194\setup.py", line 33, in <module>
          _long_description = f.read()
      UnicodeDecodeError: 'cp950' codec can't decode byte 0xe2 in position 1304: illegal multibyte sequence
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

Also I tried 0.4.1 and 0.3.25, both got the same error:

PS C:\tensorflow> pip install --no-cache-dir jax==0.4.1
Collecting jax==0.4.1
  Downloading jax-0.4.1.tar.gz (1.2 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.2/1.2 MB 3.6 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [6 lines of output]
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "C:\Users\Jason Wu\AppData\Local\Temp\pip-install-trctdimm\jax_3d064f6acc4d4a0b91bcebef44bb0eb7\setup.py", line 38, in <module>
          _long_description = f.read()
      UnicodeDecodeError: 'cp950' codec can't decode byte 0xe2 in position 1301: illegal multibyte sequence
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

PS C:\tensorflow> pip install --no-cache-dir jax==0.3.25
Collecting jax==0.3.25
  Downloading jax-0.3.25.tar.gz (1.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.1/1.1 MB 4.2 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [6 lines of output]
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "C:\Users\Jason Wu\AppData\Local\Temp\pip-install-h_ymbbfo\jax_9f7e4365d3a24fb4bb4e74d1b86678e6\setup.py", line 38, in <module>
          _long_description = f.read()
      UnicodeDecodeError: 'cp950' codec can't decode byte 0xe2 in position 1304: illegal multibyte sequence
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.

It's worth noting that 0.3.14 installed without issue:

PS C:\tensorflow> pip install --no-cache-dir jax==0.3.14
Collecting jax==0.3.14
  Downloading jax-0.3.14.tar.gz (990 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 990.1/990.1 kB 2.6 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... done
Requirement already satisfied: absl-py ...
...
Building wheels for collected packages: jax
  Building wheel for jax (setup.py) ... done
  Created wheel for jax: filename=jax-0.3.14-py3-none-any.whl size=1147576 sha256=ae0aadd2688c543930de9ba7e31e58b3363a8752ae25e4770941b53e1c3476d6
  Stored in directory: C:\Users\Jason Wu\AppData\Local\Temp\pip-ephem-wheel-cache-q92x7zxq\wheels\3d\22\cf\75c40ec058e8d0f9b8427cd35366e1fd618475451fc8fc6fd7
Successfully built jax
Installing collected packages: jax
Successfully installed jax-0.3.14

Since I'm building Tensorflow from source, it requires Jax > 0.3.15.

Please let me know if there is any other information that needs to be provided. thanks!

What jax/jaxlib version are you using?

jax v0.3.15

Which accelerator(s) are you using?

CPU

Additional system info

Python 3.9.13, pip 22.3.1, Windows 10,

NVIDIA GPU info

No response

[jakevdp]

The issue is that Windows default string encoding varies by platform. We should be able to fix this by specifying utf8 explicitly in the setup.py file, but until then you should be able to run pip install if you first set your local string encoding via an environment variable:

export LC_ALL="en_US.UTF-8"

(see https://stackoverflow.com/questions/25036897/pip-install-unicodedecodeerror)

Note, however, that the JAX team does not distribute jaxlib wheels for Windows, so even if you pip install jax you'll need to either compile jaxlib or find another source for them if you want to use JAX on Windows. See https://github.com/google/jax#installation for more information & resources.

[Jasonnor]

@jakevdp

It is troublesome to modify the local string encoding on Windows, so I installed the merged version of #13884 with pip install git+https://github.com/google/jax and it has been successfully installed.

Thanks for your help!