Math: use gather_nd batch_dims in TF >= 1.14
@@ -2,6 +2,7 @@ import uuid
import numpy as np
import six
import tensorflow as tf
+from packaging import version
from phi.math.base import Backend
from phi.struct.tensorop import expand, collapsed_gather_nd
@@ -402,9 +403,15 @@ def _resample_linear_niftynet(inputs, sample_coords, boundary, boundary_func):
def get_knot(bc):
coord = [sc[c][i] for i, c in enumerate(bc)]
+ if version.parse(tf.__version__) >= version.parse('1.14.0'):
+ coord = tf.stack(coord, -1)
+ return tf.gather_nd(inputs, coord, batch_dims=1)
+ else:
coord = tf.stack([batch_ids] + coord, -1)
return tf.gather_nd(inputs, coord)
+
+
samples = [get_knot(bc) for bc in binary_neighbour_ids]
def _pyramid_combination(samples, w_0, w_1):
