Running Jax with cudnn provided by lambda-stack

you should find it under:

/usr/lib/python3/dist-packages/tensorflow/

make sure to add the path to LD_LIBRARY_PATH variable which is for runtime and LIBRARY_PATH variable for compile time linking in your bash profile or bashrc file.