Windows & CUDA build issue

[khemritolya]

I hope whoever is reading this has had a good day. I'm attempting to build JAX on Windows with CUDA support. I'm using the following command:

python .\build\build.py --enable_cuda --cuda_version="11.7" --cudnn_version="8.4.1" --cudnn_path=C:\\Users\\.......\\cudnn-windows-x86_64-8.4.1.50_cuda11.6-archive\\cudnn-windows-x86_64-8.4.1.50_cuda11.6-archive --cuda_path="C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v11.7"

I'm using the latest version of the jax repository. CUDA 11.7, cuddn 8.4.1. I'm using bazel version 5.2.0

Here's the build error:

ERROR: @org_tensorflow//tensorflow/compiler/xla/python:enable_gpu :: Error loading option @org_tensorflow//tensorflow/compiler/xla/python:enable_gpu: no such package '@local_config_cuda//cuda': in call to path(), parameter 'path' got value of type 'NoneType', want 'string, Label, or path'

It seems to be a problem similar to one when building tensorflow? tensorflow/tensorflow#11859

Any ideas? Thanks in advance.

[hawkinsp]

JAX on Windows is supported by the community. Contributions welcome!

[mattjj]

I noticed from #5795 that there's a jax-windows-builder repo. Maybe that could help?

I hope you're having a good day too!

[hawkinsp]

Closing since I suspect the issue is stale.