Will alpa support jax 0.4.x and cuda 12.x?

alpa-projects / alpa

Training and serving large-scale neural networks with auto parallelization.

https://alpa.ai

Apache License 2.0

3.07k stars 357 forks source link

Will alpa support jax 0.4.x and cuda 12.x? #949

Open horse6 opened 1 year ago

horse6 commented 1 year ago

I tried to buid alpa using cuda 12.0, but it seems that tensorflow-alpa doesn't support it, and jax version is now at 0.4.14. Does alpa have a plan to support the newer versions of jax and cuda?