Fixes for modern JAX: block until CUDA operations complete. Import ja…

…x.experimental.
flatironinstitute · Oct 24, 2023 · aa1a634 · aa1a634
1 parent 1ca71ee
commit aa1a634
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/lib/kernels.cc.cu b/lib/kernels.cc.cu
@@ -32,6 +32,9 @@ void run_nufft(int type, const NufftDescriptor<T>* descriptor, T *x, T *y, T *z,
   }
   destroy<T>(plan);
   delete opts;
+
+  // Since JAX 0.4.9, the JAX CUDA stream is non-blocking. Need to wait for results.
+  cudaStreamSynchronize(stream);
 }
 
 

diff --git a/tests/ops_test.py b/tests/ops_test.py
@@ -2,6 +2,7 @@
 from itertools import product
 
 import jax
+import jax.experimental
 import jax.numpy as jnp
 import numpy as np
 import pytest