diff --git a/test/test_VanKeken.jl b/test/test_VanKeken.jl
index bf977d2d..c72b6783 100644
--- a/test/test_VanKeken.jl
+++ b/test/test_VanKeken.jl
@@ -1,19 +1,48 @@
 push!(LOAD_PATH, "..")
 using Test, Suppressor
 
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 2)
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 2)
+end
 
 using Printf, LinearAlgebra, GeoParams, CellArrays
 using JustRelax, JustRelax.JustRelax2D
 import JustRelax.@cell
-const backend_JR = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+
+const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
 
 using JustPIC, JustPIC._2D
 # Threads is the default backend,
 # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script,
 # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script.
-const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+const backend = @static if ENV["JULIA_JUSTPIC_BACKEND"] === "AMDGPU"
+    JustPIC.AMDGPUBackend
+elseif ENV["JULIA_JUSTPIC_BACKEND"] === "CUDA"
+    JustPIC.CUDABackend
+else
+    JustPIC.CPUBackend
+end
 
 # x-length of the domain
 const λ = 0.9142
diff --git a/test/test_diffusion3D.jl b/test/test_diffusion3D.jl
index 92126ed7..4dcf0093 100644
--- a/test/test_diffusion3D.jl
+++ b/test/test_diffusion3D.jl
@@ -1,12 +1,35 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 using GeoParams
 using JustRelax, JustRelax.JustRelax3D
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 3)
 
-const backend = CPUBackend
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 3)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 3)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 3)
+end
+
+const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
 
 # HELPER FUNCTIONS ---------------------------------------------------------------
 @parallel_indices (i, j, k) function init_T!(T, z)
@@ -110,8 +133,8 @@ function diffusion_3D(;
             args,
             dt,
             di;
-            kwargs = (; 
-                igg, 
+            kwargs = (;
+                igg,
                 verbose=false
             ),
         )
diff --git a/test/test_grid2D.jl b/test/test_grid2D.jl
index 314420cd..50058b5d 100644
--- a/test/test_grid2D.jl
+++ b/test/test_grid2D.jl
@@ -1,3 +1,11 @@
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor, JustRelax, JustRelax.JustRelax2D
 
 @testset "Grid2D" begin
diff --git a/test/test_grid3D.jl b/test/test_grid3D.jl
index 044f0ace..47359490 100644
--- a/test/test_grid3D.jl
+++ b/test/test_grid3D.jl
@@ -1,3 +1,11 @@
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor, JustRelax, JustRelax.JustRelax3D
 
 @testset "Grid3D" begin
diff --git a/test/test_shearband2D.jl b/test/test_shearband2D.jl
index bd372c4b..3e1b1e06 100644
--- a/test/test_shearband2D.jl
+++ b/test/test_shearband2D.jl
@@ -1,12 +1,35 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 using GeoParams
 using JustRelax, JustRelax.JustRelax2D
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 2)
 
-const backend = CPUBackend
+const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
+
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 2)
+end
 
 # HELPER FUNCTIONS ----------------------------------- ----------------------------
 solution(ε, t, G, η) = 2 * ε * η * (1 - exp(-G * t / η))
diff --git a/test/test_shearband2D_softening.jl b/test/test_shearband2D_softening.jl
index a8380598..3eb506b3 100644
--- a/test/test_shearband2D_softening.jl
+++ b/test/test_shearband2D_softening.jl
@@ -1,12 +1,35 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 using GeoParams, CellArrays
 using JustRelax, JustRelax.JustRelax2D
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 2)
 
-const backend = CPUBackend
+const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
+
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 2)
+end
 
 
 # HELPER FUNCTIONS ----------------------------------- ----------------------------
diff --git a/test/test_shearheating2D.jl b/test/test_shearheating2D.jl
index 06c757f0..4d1d42aa 100644
--- a/test/test_shearheating2D.jl
+++ b/test/test_shearheating2D.jl
@@ -1,19 +1,47 @@
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 
 # Benchmark of Duretz et al. 2014
 # http://dx.doi.org/10.1002/2014GL060438
 using JustRelax, JustRelax.JustRelax2D
-const backend_JR = CPUBackend
 
-using ParallelStencil, ParallelStencil.FiniteDifferences2D
-@init_parallel_stencil(Threads, Float64, 2) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2)
+const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
+
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(Threads, Float64, 2)
+end
 
 using JustPIC, JustPIC._2D
 # Threads is the default backend,
 # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script,
 # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script.
-const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
-# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+const backend = @static if ENV["JULIA_JUSTPIC_BACKEND"] === "AMDGPU"
+    JustPIC.AMDGPUBackend
+elseif ENV["JULIA_JUSTPIC_BACKEND"] === "CUDA"
+    JustPIC.CUDABackend
+else
+    JustPIC.CPUBackend
+end
 
 # Load script dependencies
 using GeoParams
diff --git a/test/test_shearheating3D.jl b/test/test_shearheating3D.jl
index ef88d44f..5854de51 100644
--- a/test/test_shearheating3D.jl
+++ b/test/test_shearheating3D.jl
@@ -1,21 +1,52 @@
 push!(LOAD_PATH, "..")
 using Test, Suppressor
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 # Benchmark of Duretz et al. 2014
 # http://dx.doi.org/10.1002/2014GL060438
 using JustRelax, JustRelax.JustRelax3D
-const backend_JR = CPUBackend
-import JustRelax.@cell
 
-using ParallelStencil, ParallelStencil.FiniteDifferences3D
-@init_parallel_stencil(Threads, Float64, 3) #or (CUDA, Float64, 2) or (AMDGPU, Float64, 2)
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 3)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 3)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 3)
+end
+
 
-using JustPIC, JustPIC._3D
+const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
+
+
+using JustPIC
+using JustPIC._3D
 # Threads is the default backend,
 # to run on a CUDA GPU load CUDA.jl (i.e. "using CUDA") at the beginning of the script,
 # and to run on an AMD GPU load AMDGPU.jl (i.e. "using AMDGPU") at the beginning of the script.
-const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
-# const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+const backend = @static if ENV["JULIA_JUSTPIC_BACKEND"] === "AMDGPU"
+    JustPIC.AMDGPUBackend
+elseif ENV["JULIA_JUSTPIC_BACKEND"] === "CUDA"
+    CUDABackend
+else
+    JustPIC.CPUBackend
+end
+import JustRelax.@cell
 
 # Load script dependencies
 using Printf, GeoParams
diff --git a/test/test_sinking_block.jl b/test/test_sinking_block.jl
index 927060f1..1f8f0728 100644
--- a/test/test_sinking_block.jl
+++ b/test/test_sinking_block.jl
@@ -1,14 +1,44 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
+
 using Test, Suppressor
 using JustRelax, JustRelax.JustRelax2D
-using ParallelStencil, ParallelStencil.FiniteDifferences2D
-@init_parallel_stencil(Threads, Float64, 2)
 
-const backend_JR = CPUBackend
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(Threads, Float64, 2)
+end
+
+const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
 
 using JustPIC, JustPIC._2D
-const backend = CPUBackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
+const backend = @static if ENV["JULIA_JUSTPIC_BACKEND"] === "AMDGPU"
+    JustPIC.AMDGPUBackend
+elseif ENV["JULIA_JUSTPIC_BACKEND"] === "CUDA"
+    CUDABackend
+else
+    JustPIC.CPUBackend
+end
 
 using GeoParams
 
diff --git a/test/test_stokes_burstedde.jl b/test/test_stokes_burstedde.jl
index fbccbbeb..597d062e 100644
--- a/test/test_stokes_burstedde.jl
+++ b/test/test_stokes_burstedde.jl
@@ -1,13 +1,26 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
-using JustRelax
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 3)
+using JustRelax, JustRelax.JustRelax3D
 
-# setup ParallelStencil.jl environment
-model = PS_Setup(:cpu, Float64, 3)
-environment!(model)
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 3)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 3)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 3)
+end
 
 include("../miniapps/benchmarks/stokes3D/burstedde/Burstedde.jl")
 
diff --git a/test/test_stokes_elastic_buildup.jl b/test/test_stokes_elastic_buildup.jl
index 7ae7e950..6b2c9aee 100644
--- a/test/test_stokes_elastic_buildup.jl
+++ b/test/test_stokes_elastic_buildup.jl
@@ -1,11 +1,34 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 using JustRelax, JustRelax.JustRelax2D
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 2)
 
-const backend = CPUBackend
+const backend_JR = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
+
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil, ParallelStencil.FiniteDifferences2D
+    @init_parallel_stencil(Threads, Float64, 2)
+end
 
 include("../miniapps/benchmarks/stokes2D/elastic_buildup/Elastic_BuildUp.jl")
 
diff --git a/test/test_stokes_solcx.jl b/test/test_stokes_solcx.jl
index a4034172..a1d68821 100644
--- a/test/test_stokes_solcx.jl
+++ b/test/test_stokes_solcx.jl
@@ -1,11 +1,34 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 using JustRelax, JustRelax.JustRelax2D
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 2)
 
-const backend = CPUBackend
+const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
+
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 2)
+end
 
 include("../miniapps/benchmarks/stokes2D/solcx/SolCx.jl")
 
diff --git a/test/test_stokes_solkz.jl b/test/test_stokes_solkz.jl
index 9c0246c1..f716d4fa 100644
--- a/test/test_stokes_solkz.jl
+++ b/test/test_stokes_solkz.jl
@@ -1,11 +1,35 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 using JustRelax, JustRelax.JustRelax2D
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 2)
 
-const backend = CPUBackend
+const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
+
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 2)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 2)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 2)
+end
+
 
 include("../miniapps/benchmarks/stokes2D/solkz/SolKz.jl")
 
diff --git a/test/test_stokes_solvi3D.jl b/test/test_stokes_solvi3D.jl
index 170eda32..061a8b46 100644
--- a/test/test_stokes_solvi3D.jl
+++ b/test/test_stokes_solvi3D.jl
@@ -1,11 +1,34 @@
 push!(LOAD_PATH, "..")
 
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using Test, Suppressor
 using JustRelax, JustRelax.JustRelax3D
-const backend = CPUBackend
 
-using ParallelStencil
-@init_parallel_stencil(Threads, Float64, 3)
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using ParallelStencil
+    @init_parallel_stencil(AMDGPU, Float64, 3)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using ParallelStencil
+    @init_parallel_stencil(CUDA, Float64, 3)
+else
+    using ParallelStencil
+    @init_parallel_stencil(Threads, Float64, 3)
+end
+
+const backend = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
 
 include("../miniapps/benchmarks/stokes3D/solvi/SolVi3D.jl")
 
diff --git a/test/test_traits.jl b/test/test_traits.jl
index bc942702..bc7fda23 100644
--- a/test/test_traits.jl
+++ b/test/test_traits.jl
@@ -1,44 +1,135 @@
+@static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    using AMDGPU
+    AMDGPU.allowscalar(true)
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    using CUDA
+    CUDA.allowscalar(true)
+end
+
 using JustRelax, Test
 import JustRelax.JustRelax2D as JR2
 import JustRelax.JustRelax3D as JR3
 
-const bk = JustRelax.backend
+const bk = @static if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+    JustRelax.AMDGPUBackend
+elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+    JustRelax.CUDABackend
+else
+    JustRelax.CPUbackend
+end
 
 @testset "Traits" begin
-    # test generic arrays
-    @test bk(Array)       === CPUBackendTrait()
-    @test bk(Matrix)      === CPUBackendTrait()
-    @test bk(Vector)      === CPUBackendTrait()
-    @test bk(rand(2))     === CPUBackendTrait()
-    @test bk(rand(2,2))   === CPUBackendTrait()
-    @test bk(rand(2,2,2)) === CPUBackendTrait()
-
-    # test error handling
-    @test_throws ArgumentError bk(rand()) 
-    @test_throws ArgumentError bk("potato")
-
-    # test JR structs
-    ## 2D
-    ni       = 2, 2
-    stokes2  = JR2.StokesArrays(CPUBackend, ni)
-    thermal2 = JR2.ThermalArrays(CPUBackend, ni)
-
-    @test bk(stokes2.V) === CPUBackendTrait()
-    @test bk(stokes2.τ) === CPUBackendTrait()
-    @test bk(stokes2.R) === CPUBackendTrait()
-    @test bk(stokes2.P) === CPUBackendTrait()
-    @test bk(stokes2)   === CPUBackendTrait()
-    @test bk(thermal2)  === CPUBackendTrait()
-    
-    ## 3D
-    ni       = 2, 2, 2
-    stokes3  = JR3.StokesArrays(CPUBackend, ni)
-    thermal3 = JR3.ThermalArrays(CPUBackend, ni)
- 
-    @test bk(stokes3.V) === CPUBackendTrait()
-    @test bk(stokes3.τ) === CPUBackendTrait()
-    @test bk(stokes3.R) === CPUBackendTrait()
-    @test bk(stokes3.P) === CPUBackendTrait()
-    @test bk(stokes3)   === CPUBackendTrait()
-    @test bk(thermal3)  === CPUBackendTrait()
-end
\ No newline at end of file
+    if ENV["JULIA_JUSTRELAX_BACKEND"] === "AMDGPU"
+        # test generic arrays
+        @test bk(Array)       === AMDGPUBackendTrait()
+        @test bk(Matrix)      === AMDGPUBackendTrait()
+        @test bk(Vector)      === AMDGPUBackendTrait()
+        @test bk(rand(2))     === AMDGPUBackendTrait()
+        @test bk(rand(2,2))   === AMDGPUBackendTrait()
+        @test bk(rand(2,2,2)) === AMDGPUBackendTrait()
+
+        # test error handling
+        @test_throws ArgumentError bk(rand())
+        @test_throws ArgumentError bk("potato")
+
+        # test JR structs
+        ## 2D
+        ni       = 2, 2
+        stokes2  = JR2.StokesArrays(AMDGPUBackend, ni)
+        thermal2 = JR2.ThermalArrays(AMDGPUBackend, ni)
+
+        @test bk(stokes2.V) === AMDGPUBackendTrait()
+        @test bk(stokes2.τ) === AMDGPUBackendTrait()
+        @test bk(stokes2.R) === AMDGPUBackendTrait()
+        @test bk(stokes2.P) === AMDGPUBackendTrait()
+        @test bk(stokes2)   === AMDGPUBackendTrait()
+        @test bk(thermal2)  === AMDGPUBackendTrait()
+
+        ## 3D
+        ni       = 2, 2, 2
+        stokes3  = JR3.StokesArrays(AMDGPUBackend, ni)
+        thermal3 = JR3.ThermalArrays(AMDGPUBackend, ni)
+
+        @test bk(stokes3.V) === AMDGPUBackendTrait()
+        @test bk(stokes3.τ) === AMDGPUBackendTrait()
+        @test bk(stokes3.R) === AMDGPUBackendTrait()
+        @test bk(stokes3.P) === AMDGPUBackendTrait()
+        @test bk(stokes3)   === AMDGPUBackendTrait()
+        @test bk(thermal3)  === AMDGPUBackendTrait()
+    elseif ENV["JULIA_JUSTRELAX_BACKEND"] === "CUDA"
+        # test generic arrays
+        @test bk(Array)       === CUDABackendTrait()
+        @test bk(Matrix)      === CUDABackendTrait()
+        @test bk(Vector)      === CUDABackendTrait()
+        @test bk(rand(2))     === CUDABackendTrait()
+        @test bk(rand(2,2))   === CUDABackendTrait()
+        @test bk(rand(2,2,2)) === CUDABackendTrait()
+
+        # test error handling
+        @test_throws ArgumentError bk(rand())
+        @test_throws ArgumentError bk("potato")
+
+        # test JR structs
+        ## 2D
+        ni       = 2, 2
+        stokes2  = JR2.StokesArrays(CUDABackend, ni)
+        thermal2 = JR2.ThermalArrays(CUDABackend, ni)
+
+        @test bk(stokes2.V) === CUDABackendTrait()
+        @test bk(stokes2.τ) === CUDABackendTrait()
+        @test bk(stokes2.R) === CUDABackendTrait()
+        @test bk(stokes2.P) === CUDABackendTrait()
+        @test bk(stokes2)   === CUDABackendTrait()
+        @test bk(thermal2)  === CUDABackendTrait()
+
+        ## 3D
+        ni       = 2, 2, 2
+        stokes3  = JR3.StokesArrays(CUDABackend, ni)
+        thermal3 = JR3.ThermalArrays(CUDABackend, ni)
+
+        @test bk(stokes3.V) === CUDABackendTrait()
+        @test bk(stokes3.τ) === CUDABackendTrait()
+        @test bk(stokes3.R) === CUDABackendTrait()
+        @test bk(stokes3.P) === CUDABackendTrait()
+        @test bk(stokes3)   === CUDABackendTrait()
+        @test bk(thermal3)  === CUDABackendTrait()
+    else
+        # test generic arrays
+        @test bk(Array)       === CPUBackendTrait()
+        @test bk(Matrix)      === CPUBackendTrait()
+        @test bk(Vector)      === CPUBackendTrait()
+        @test bk(rand(2))     === CPUBackendTrait()
+        @test bk(rand(2,2))   === CPUBackendTrait()
+        @test bk(rand(2,2,2)) === CPUBackendTrait()
+
+        # test error handling
+        @test_throws ArgumentError bk(rand())
+        @test_throws ArgumentError bk("potato")
+
+        # test JR structs
+        ## 2D
+        ni       = 2, 2
+        stokes2  = JR2.StokesArrays(CPUBackend, ni)
+        thermal2 = JR2.ThermalArrays(CPUBackend, ni)
+
+        @test bk(stokes2.V) === CPUBackendTrait()
+        @test bk(stokes2.τ) === CPUBackendTrait()
+        @test bk(stokes2.R) === CPUBackendTrait()
+        @test bk(stokes2.P) === CPUBackendTrait()
+        @test bk(stokes2)   === CPUBackendTrait()
+        @test bk(thermal2)  === CPUBackendTrait()
+
+        ## 3D
+        ni       = 2, 2, 2
+        stokes3  = JR3.StokesArrays(CPUBackend, ni)
+        thermal3 = JR3.ThermalArrays(CPUBackend, ni)
+
+        @test bk(stokes3.V) === CPUBackendTrait()
+        @test bk(stokes3.τ) === CPUBackendTrait()
+        @test bk(stokes3.R) === CPUBackendTrait()
+        @test bk(stokes3.P) === CPUBackendTrait()
+        @test bk(stokes3)   === CPUBackendTrait()
+        @test bk(thermal3)  === CPUBackendTrait()
+    end
+
+end