tweak cuda 12/13

Author: cjerzak

fastrerandomize/R/FRR_BuildBackend.R

@@ -25,7 +25,7 @@ build_backend <- function(conda_env = "fastrerandomize_env", conda = "auto"){
   # Create a new conda environment
   reticulate::conda_create(envname = conda_env,
                            conda = conda,
-                           python_version = "3.11")
+                           python_version = "3.12")
 
   os <- Sys.info()[["sysname"]]
   machine <- Sys.info()["machine"]
@@ -52,20 +52,20 @@ build_backend <- function(conda_env = "fastrerandomize_env", conda = "auto"){
       # Prefer CUDA 13 if the driver is new enough; otherwise CUDA 12; else CPU fallback
       if (!is.na(drv_major) && drv_major >= 580) {
         msg("Driver %s detected (>=580): installing JAX 0.5.0 CUDA 13 wheels.", drv[1])
-        tryCatch(pip_install('jax[cuda13]==0.5.0'), error = function(e) {
+        tryCatch(pip_install('jax[cuda13]'), error = function(e) {
           msg("CUDA 13 wheels failed (%s); falling back to CUDA 12.", e$message)
-          pip_install('jax[cuda12]==0.5.0')
+          pip_install('jax[cuda12]')
         })
       } else if (!is.na(drv_major) && drv_major >= 525) {
         msg("Driver %s detected (>=525,<580): installing JAX 0.5.0 CUDA 12 wheels.", drv[1])
-        pip_install('jax[cuda12]==0.5.0')
+        pip_install('jax[cuda12]')
       } else {
         msg("Driver %s too old for CUDA wheels; installing CPU-only JAX 0.5.0.", drv[1])
-        pip_install('jax==0.5.0')
+        pip_install('jax')
       }
     } else {
       msg("Installing CPU-only JAX 0.5.0.")
-      pip_install('jax==0.5.0')
+      pip_install('jax')
     }
   }