[DRAFT] Jax

[johnnynunez]

TIPS:

1. cudnn8 is not compatible with next releases, so maybe we have to upgrade 12.2 image to 9.0 version. Or developers must be do a rev 3 of jetpack 6.
2. cudnn9 must be renamed and mantain a specific structure of folders. (I did a script that do that in /opt/nvidia/cudnn)
3. clang not detect cuda and cudnn, so it must be passes by command.
4. create_cudnn is like a temporal folder for build jax.
5. All must be match if not, crash!

IDEA:
Create a docker.builder that run create_cudnn.sh etc and build jax. All must be automatize by hyperparameter, like sm_87, cuda version, cudnn_version and local_cuda_path

New command:

python3 build/build.py --enable_cuda --enable_nccl=False \
        --cuda_compute_capabilities="sm_87" \
        --cuda_version=12.2.0 --cudnn_version=9.4.0 \
        --bazel_options=--repo_env=LOCAL_CUDA_PATH="/usr/local/cuda-12.2" \
        --bazel_options=--repo_env=LOCAL_CUDNN_PATH="/opt/nvidia/cudnn/"

[johnnynunez]

#618

[johnnynunez]

@dusty-nv all is prepared for the next release 0.4.32.
Now depends on the commit builds or not. In my case, is building on my jetson outsde dockers. So I will wait for the next release in the coming days.
Dockers are not prepared. taking accont my tips

[johnnynunez]

I've improved the script because needs that cuda must be lib instead of lib64 path