updated for tf >= 2.9

.github/workflows/test.yml

@@ -17,10 +17,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.8
+    - name: Set up Python 3.9
       uses: actions/setup-python@v2
       with:
-        python-version: 3.8
+        python-version: 3.9
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

requirements.txt

@@ -1,3 +1,3 @@
-tensorflow>=2.1.0
-numpy
+tensorflow>=2.9.0
+numpy <2.0 
 scipy

setup.py

@@ -68,6 +68,6 @@ def load_requirements(path_dir=HERE, comment_char='#'):
         # Specify the Python versions you support here. In particular, ensure
         # that you indicate whether you support Python 2, Python 3 or both.
         'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.9',
     ],
 )

tfkbnufft/kbnufft.py

@@ -150,7 +150,7 @@ def _extract_nufft_interpob(self):
         return interpob
 
 def kbnufft_forward(interpob, multiprocessing=False):
-    @tf.function(experimental_relax_shapes=True)
+    @tf.function(reduce_retracing=True)
     @tf.custom_gradient
     def kbnufft_forward_for_interpob(x, om):
         """Apply FFT and interpolate from gridded data to scattered data.
@@ -207,7 +207,7 @@ def grad(dy):
     return kbnufft_forward_for_interpob
 
 def kbnufft_adjoint(interpob, multiprocessing=False):
-    @tf.function(experimental_relax_shapes=True)
+    @tf.function(reduce_retracing=True)
     @tf.custom_gradient
     def kbnufft_adjoint_for_interpob(y, om):
         """Interpolate from scattered data to gridded data and then iFFT.

tfkbnufft/nufft/interp_functions.py

@@ -57,7 +57,7 @@ def calc_coef_and_indices(tm, kofflist, Jval, table, centers, L, dims, conjcoef=
 
     return coef, arr_ind
 
-@tf.function(experimental_relax_shapes=True)
+@tf.function(reduce_retracing=True)
 def run_interp(griddat, tm, params):
     """Interpolates griddat to off-grid coordinates with input sparse matrices.
 
@@ -107,7 +107,7 @@ def run_interp(griddat, tm, params):
 
     return kdat
 
-@tf.function(experimental_relax_shapes=True)
+@tf.function(reduce_retracing=True)
 def run_interp_back(kdat, tm, params):
     """Interpolates kdat to on-grid coordinates.
 
@@ -163,7 +163,7 @@ def run_interp_back(kdat, tm, params):
     griddat = tf.transpose(tf.complex(griddat_real, griddat_imag))
     return griddat
 
-@tf.function(experimental_relax_shapes=True)
+@tf.function(reduce_retracing=True)
 def kbinterp(x, om, interpob, conj=False):
     """Apply table interpolation.
 
@@ -219,11 +219,11 @@ def _map_body(inputs):
             y = y_not_shifted * shift
         return y
 
-    y = tf.map_fn(_map_body, [x, tm, om], dtype=x.dtype)
+    y = tf.map_fn(_map_body, [x, tm, om], fn_output_signature=x.dtype)
 
     return y
 
-@tf.function(experimental_relax_shapes=True)
+@tf.function(reduce_retracing=True)
 def adjkbinterp(y, om, interpob):
     """Apply table interpolation adjoint.
 
@@ -272,7 +272,7 @@ def _map_body(inputs):
         x = run_interp_back(y_shifted, _tm, params)
         return x
 
-    x = tf.map_fn(_map_body, [y, om, tm], dtype=y.dtype)
+    x = tf.map_fn(_map_body, [y, om, tm], fn_output_signature=y.dtype)
 
     bsize = tf.shape(y)[0]
     ncoil = tf.shape(y)[1]