"Could not find any cudnn.h" when building from source

Starting on commit https://github.com/google/jax/commit/60d7198381db2098c3319c01ba8060bad9a14422 I get the following error when building with cuda support. Commits before that build correctly.

$ python build/build.py --enable_cuda

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Extracting Bazel installation...
Starting local Bazel server and connecting to it...
Bazel binary path: ./bazel-0.24.0-linux-x86_64
Python binary path: /opt/conda/bin/python
MKL-DNN enabled: yes
-march=native: no
CUDA enabled: yes
CUDA toolkit path: /usr/local/cuda
CUDNN library path: /usr/local/cuda

Building XLA and installing it in the jaxlib source tree...
ERROR: /root/code/jax/build/BUILD.bazel:21:1: error loading package '@org_tensorflow//tensorflow/compiler/xla/python': in /root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/tensorflow/compiler/xla/xla.bzl: Encountered error while reading extension file 'cuda/build_defs.bzl': no such package '@local_config_cuda//cuda': Traceback (most recent call last):
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 1253
		_create_local_cuda_repository(repository_ctx)
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 956, in _create_local_cuda_repository
		_get_cuda_config(repository_ctx)
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 688, in _get_cuda_config
		find_cuda_config(repository_ctx, ["cuda", "cudnn"])
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 668, in find_cuda_config
		auto_configure_fail(("Failed to run find_cuda_config...))
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 274, in auto_configure_fail
		fail(("\n%sCuda Configuration Error:%...)))

Cuda Configuration Error: Failed to run find_cuda_config.py: Could not find any cudnn.h matching version '' in any subdirectory:
        ''
        'include'
        'include/cuda'
        'include/*-linux-gnu'
        'extras/CUPTI/include'
        'include/cuda/CUPTI'
of:
        '/usr/local/cuda'

 and referenced by '//build:install_xla_in_source_tree'
ERROR: Analysis of target '//build:install_xla_in_source_tree' failed; build aborted: error loading package '@org_tensorflow//tensorflow/compiler/xla/python': in /root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/tensorflow/compiler/xla/xla.bzl: Encountered error while reading extension file 'cuda/build_defs.bzl': no such package '@local_config_cuda//cuda': Traceback (most recent call last):
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 1253
		_create_local_cuda_repository(repository_ctx)
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 956, in _create_local_cuda_repository
		_get_cuda_config(repository_ctx)
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 688, in _get_cuda_config
		find_cuda_config(repository_ctx, ["cuda", "cudnn"])
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 668, in find_cuda_config
		auto_configure_fail(("Failed to run find_cuda_config...))
	File "/root/.cache/bazel/_bazel_root/180cc62733246893a6f39cd7c8b63394/external/org_tensorflow/third_party/gpus/cuda_configure.bzl", line 274, in auto_configure_fail
		fail(("\n%sCuda Configuration Error:%...)))

Cuda Configuration Error: Failed to run find_cuda_config.py: Could not find any cudnn.h matching version '' in any subdirectory:
        ''
        'include'
        'include/cuda'
        'include/*-linux-gnu'
        'extras/CUPTI/include'
        'include/cuda/CUPTI'
of:
        '/usr/local/cuda'

[hawkinsp]

I think TensorFlow changed how it finds CUDA.

Try this patch. It removes the --cudnn_path option and lets TF’s logic find CuDNN automatically. https://gist.github.com/hawkinsp/29139c1edfa60e7ddaea8143fdcbcb5e

(In general our jaxlib" build does unusual things with the TF build process to build XLA. We should probably move the jaxlib build entirely into the XLA tree so we aren’t as prone to this kind of breakage. jaxlib is really xlapy anyway.)

[hawkinsp]

Please open a new issue. I don’t think this is the same problem. a) if you haven’t considered it, you should prefer to use a prebuilt jaxlib. b) everything that patch did is present in master (i.e., --cudnn_path is optional, and if you don’t specify it, it acts the same way as that patch did.)