Example in Deep Learning with R, 3rd edition does not work with JAX backend on Linux / NVIDIA GPU

[znmeb]

I'm working my way through Deep Learning with R, 3rd Edition. In chapter 3 there are three examples. The examples with the TensorFlow and PyTorch back ends work fine, but the JAX one gets what looks like some kind of error in JAX trying to access the GPU:

#! /usr/bin/env Rscript

update.packages(ask = FALSE, repos = "https://cloud.r-project.org/")

required_packages <- c(
  "keras3",
  "tensorflow"
)
install.packages(required_packages, quiet = TRUE, repos = "https://cloud.r-project.org/")
warnings()

# Deep Learning with R, 3rd edition, Chapter section 3.5.1
library(keras3)
use_backend("jax")
jax <- import("jax")

# Deep Learning with R, 3rd edition, Chapter section 3.5.2
jnp <- import("jax.numpy")
jnp$ones(shape = shape(2, 2))
warnings()

gives

TEP 5/7: RUN nvidia-smi
Thu Jul 30 01:18:37 2026       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 610.43.03              KMD Version: 610.43.03     CUDA UMD Version: 13.3     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3090        Off |   00000000:01:00.0  On |                  N/A |
| 32%   34C    P8             14W /  350W |     993MiB /  24576MiB |     15%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|  No running processes found                                                             |
+-----------------------------------------------------------------------------------------+
STEP 6/7: COPY R-script.R ./
STEP 7/7: RUN ./R-script.R
also installing the dependencies 'lattice', 'rprojroot', 'vctrs', 'Matrix', 'Rcpp', 'RcppTOML', 'here', 'jsonlite', 'png', 'rappdirs', 'withr', '
base64enc', 'whisker', 'tidyselect', 'ps', 'R6', 'backports', 'generics', 'reticulate', 'tfruns', 'magrittr', 'zeallot', 'fastmap', 'glue', 'cli'
, 'rlang', 'dotty', 'config', 'processx', 'yaml', 'tfautograph', 'rstudioapi', 'lifecycle'

Downloading uv...Done!
Downloading cpython-3.12.13-linux-x86_64-gnu (download) (32.6MiB)
 Downloaded cpython-3.12.13-linux-x86_64-gnu (download)
Downloading jaxlib (83.2MiB)
Downloading grpcio (6.7MiB)
Downloading keras (2.3MiB)
Downloading numpy (15.9MiB)
Downloading jax-cuda12-plugin (7.8MiB)
Downloading nvidia-cublas-cu12 (554.3MiB)
Downloading nvidia-nvjitlink-cu12 (37.9MiB)
Downloading nvidia-cuda-runtime-cu12 (3.3MiB)
Downloading nvidia-cuda-nvrtc-cu12 (85.4MiB)
Downloading tensorflow-cpu (261.3MiB)
Downloading jax (3.1MiB)
Downloading pygments (1.2MiB)
Downloading pillow (6.6MiB)
Downloading nvidia-cufft-cu12 (191.6MiB)
Downloading jax-cuda12-pjrt (167.7MiB)
Downloading jedi (4.7MiB)
Downloading ml-dtypes (4.8MiB)
Downloading nvidia-nccl-cu12 (289.3MiB)
Downloading nvidia-cuda-cccl-cu12 (3.0MiB)
Downloading pandas (10.5MiB)
Downloading nvidia-nvshmem-cu12 (219.5MiB)
Downloading nvidia-cusparse-cu12 (349.5MiB)
Downloading scipy (33.7MiB)
Downloading nvidia-cusolver-cu12 (322.5MiB)
Downloading libclang (23.4MiB)
Downloading nvidia-cuda-cupti-cu12 (10.3MiB)
Downloading h5py (4.7MiB)
Downloading nvidia-cuda-nvcc-cu12 (38.7MiB)
Downloading nvidia-cudnn-cu12 (762.1MiB)
 Downloaded pygments
 Downloaded keras
 Downloaded nvidia-cuda-cccl-cu12
 Downloaded jax
 Downloaded nvidia-cuda-runtime-cu12
 Downloaded h5py
 Downloaded jedi
 Downloaded ml-dtypes
 Downloaded grpcio
 Downloaded pillow
 Downloaded jax-cuda12-plugin
 Downloaded nvidia-cuda-cupti-cu12
 Downloaded pandas
 Downloaded numpy
 Downloaded libclang
 Downloaded scipy
 Downloaded nvidia-nvjitlink-cu12
 Downloaded nvidia-cuda-nvcc-cu12
 Downloaded jaxlib
 Downloaded nvidia-cuda-nvrtc-cu12
 Downloaded jax-cuda12-pjrt
 Downloaded nvidia-cufft-cu12
 Downloaded nvidia-nvshmem-cu12
 Downloaded nvidia-nccl-cu12
 Downloaded tensorflow-cpu
 Downloaded nvidia-cusolver-cu12
 Downloaded nvidia-cusparse-cu12
 Downloaded nvidia-cublas-cu12
 Downloaded nvidia-cudnn-cu12
Installed 71 packages in 893ms
E0730 01:24:29.942174      45 cuda_executor.cc:1176] [0] Failed to allocate device memory of 23.56GiB (25293750272 bytes): RESOURCE_EXHAUSTED: : 
CUDA_ERROR_OUT_OF_MEMORY: out of memory
=== Source Location Trace: ===
external/xla/xla/stream_executor/cuda/cuda_status.cc:45
external/xla/xla/stream_executor/cuda/cuda_device_allocator.cc:226
external/xla/xla/stream_executor/cuda/cuda_device_allocator.cc:403

Array([[1., 1.],
       [1., 1.]], dtype=float32)

This looks like a JAX problem - when I watch this with nvtop I can see it allocating the memory and then de-allocating it. But I have no idea how to go about troubleshooting JAX.

All of this code is in https://github.com/AlgoCompSynth/jax-reprex. I'm running it in an Ubuntu 26.04 Podman container but it should run in any Linux machine with an NVIDIA GPU and R.

[t-kalinowski]

Thank you very much for the bug report and reproducible example! This is now fixed on main by #1539.

keras3 was setting XLA_PYTHON_CLIENT_MEM_FRACTION=1.00, causing JAX to attempt to preallocate all available GPU memory. We now leave that variable unset, preserving JAX’s default of 75% or any user-provided value.