Update Mac jaxlib?

[tbenthompson]

Hi!

I got here from this issue comment (python-poetry/poetry#6956 (comment)). It looks like the Mac version of the jaxlib hasn't been updated since 0.3.18 (https://jorenham.github.io/jax_pep503/jaxlib/). It would be wonderful to be able to use this!

Thanks!!

[jorenham]

The repo uses this (non pep503) one as source: https://storage.googleapis.com/jax-releases/
The latest mac release is 0.3.18 there as well, so I'm afraid that there isn't anything that can be done about this.

[jorenham]

There aren't even any mac-specific wheels here: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[tbenthompson]

Oh, ok. I misunderstood. Thanks regardless!

The point of the poetry issue that I got here from was to be able to have a cross-platform poetry pyproject.toml that:

1. Would install a CUDA-capable version of jax/jaxlib on Linux machines.
2. Install a non-CUDA-capable version of jax/jaxlib on a Mac.

FYI, the Mac-specific wheels are stored here: https://storage.googleapis.com/jax-releases/jax_releases.html

[tbenthompson]

I'm not sure where you are seeing that the latest release is 0.3.18. I see all the most recent releases!

[jorenham]

@tbenthompson
I think I have found the issue here.

As you can see in the code, the contents of https://storage.googleapis.com/jax-releases/ are fetched, and in there, the newer mac releases are indeed missing.

I noticed only now that at the top of the xml, there's a <NextMarker>, indicating pagination. So currently, only the first "page" of releases is indexed.

After digging a bit more, I found that the returned XML is actually a response from the AWS S3 ListObjects action. This means boto3 can be used to do the parsing.

I'll look into fixing this today or tomorrow.

And thanks @tbenthompson for pointing this out; this is a rather significant (and sneaky) bug.

[tbenthompson]

Oh that's awesome! Glad you tracked down the issue and I'm psyched to use this.

[jorenham]

@tbenthompson The fixes are live now :)

[NeilGirdhar]

I can't seem to install on Mac. Is this working for anyone? Can anyone share their your pyproject.toml dependency?

[jorenham]

@NeilGirdhar I believe you have answered this question yourself in #6 :)

[NeilGirdhar]

@jorenham I think I'm essentially running into this: python-poetry/poetry#6956 Thanks for your repo by the way 😄

[jorenham]

@NeilGirdhar I'm a linux user myself, so I'm afraid I can't be of much assistance with max related issues 🤷🏻. I assume that you've found the working jax+cuda example project already?

Otherwise, It'd be best to open a new issue, and provide the poetry output and other relevant info.

[NeilGirdhar]

@jorenham Yeah, I'm a linux user too. I just want the same pyproject to work on both Linux and Mac. I guess that's impossible for now in Poetry. Thanks for the link.

[jorenham]

@tbenthompson Did you manage to get it working in the end?

[NeilGirdhar]

Nope! I just use pip on my mac for now. I mainly develop on linux, so it's a minor inconvenience. There does seem to be some problems with using Poetry on multiple platforms though.

[jorenham]

@NeilGirdhar, my question was directed at @tbenthompson

[tbenthompson]

No, I never got the cross-platform poetry install working in a way that would install the lates CUDA version on Linux and the latest CPU version on Mac. I ended up just having poetry install the CPU version and then manually running pip to replace the CPU version with CUDA on Linux if desired.

[NeilGirdhar]

@tbenthompson Good to know. That's exactly what I'm trying to do so if you remember, please let me know if you find a way one day.