jax fixings

Dockerfile

@@ -30,10 +30,10 @@ RUN set -eux; \
 ENV LD_PRELOAD=/usr/local/cuda/lib64/libcusparse.so.12:/usr/local/cuda/lib64/libcublas.so.12:/usr/local/cuda/lib64/libcublasLt.so.12:/usr/local/cuda/lib64/libcufft.so.11:/usr/local/cuda/lib64/libcusolver.so.11
 
 # Better allocator (less fragmentation than BFC during XLA autotune)
-# ENV TF_GPU_ALLOCATOR=cuda_malloc_async
+ENV TF_GPU_ALLOCATOR=cuda_malloc_async
 
 # Let cuBLAS use TF32 fast path on Ada (L40S) for big GEMMs
-# ENV TF_ENABLE_CUBLAS_TF32=1 NVIDIA_TF32_OVERRIDE=1
+ENV TF_ENABLE_CUBLAS_TF32=1 NVIDIA_TF32_OVERRIDE=1
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
@@ -128,6 +128,9 @@ RUN python -m pip install --no-cache-dir --force-reinstall "protobuf==4.25.3"
 
 RUN python -m pip install gradio
 
+RUN python -m pip install --no-cache-dir --force-reinstall \
+    "jax==0.7.1" "jaxlib==0.7.1" "jax[cuda12]==0.7.1"
+
 # Switch to Spaces’ preferred user
 # Switch to Spaces’ preferred user
 RUN useradd -m -u 1000 appuser