Tensorflow map_fn gives error: ValueError: No attr named '_XlaCompile' - tensorflow

Question

I try to implement the 'batch hard' batches as described in https://arxiv.org/pdf/1703.07737.pdf to use with a triplet loss. So input is of shape [batch_size, 32] and output should be a list representing triplets, so [[batch_size, 32], [batch_size, 32], [batch_size, 32]] when each individual example is of size (32,).
I implemented this with the following function, so basically using tf.map_fn:
def batch_hard(inputs):
"""
Batch Hard triplets as described in https://arxiv.org/pdf/1703.07737.pdf.
For each sample in input the hardest positive and hardest negative
in the given batch will be selected. A triplet is returned.
"""
class_ids, f_anchor = inputs[0], inputs[1]
def body(x):
class_id, f = x[0], x[1]
same_class = tf.equal(class_ids, class_id)
positive = same_class
negative = tf.logical_not(same_class)
positive = tf.squeeze(positive)
negative = tf.squeeze(negative)
positive.set_shape([None])
negative.set_shape([None])
samples_pos = tf.boolean_mask(f_anchor, positive)
samples_neg = tf.boolean_mask(f_anchor, negative)
# Select hardest positive example
distances = euclidean_distance(samples_pos, f)
hardest_pos = samples_pos[tf.argmax(distances)]
# Select hardest negative example
distances = euclidean_distance(samples_neg, f)
hardest_neg = samples_neg[tf.argmin(distances)]
return [hardest_pos, hardest_neg]
[f_pos, f_neg] = tf.map_fn(body, inputs, dtype=[tf.float32, tf.float32])
return [f_anchor, f_pos, f_neg]
This works perfectly when I only perform a forward pass, with no train_op specified . However when I add this line train_op = optimizer.minimize(loss, global_step=global_step) the following error occurs:
Traceback (most recent call last):
File "/usr/local/lib/python3.5/dist-packages/tensorflow/python/ops/gradients_impl.py", line 348, in _MaybeCompile
xla_compile = op.get_attr("_XlaCompile")
File "/usr/local/lib/python3.5/dist-packages/tensorflow/python/framework/ops.py", line 2003, in get_attr
raise ValueError("No attr named '" + name + "' in " + str(self._node_def))
ValueError: No attr named '_XlaCompile' in name: "map/while/strided_slice"
op: "StridedSlice"
input: "map/while/boolean_mask/Gather"
input: "map/while/strided_slice/stack"
input: "map/while/strided_slice/stack_1"
input: "map/while/strided_slice/Cast"
attr {
key: "Index"
value {
type: DT_INT64
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "begin_mask"
value {
i: 0
}
}
attr {
key: "ellipsis_mask"
value {
i: 0
}
}
attr {
key: "end_mask"
value {
i: 0
}
}
attr {
key: "new_axis_mask"
value {
i: 0
}
}
attr {
key: "shrink_axis_mask"
value {
i: 1
}
}
Does anyone has an idea what goes wrong?
A full example of this issue is here https://gist.github.com/anonymous/0b5e9194ebf09be7ad2f0a740bf369b8
Edit: It seems the problems is in these lines
hardest_pos = samples_pos[tf.argmax(distances)]
replacing it with something like
hardest_pos = tf.zeros(32)
gives no errors, however how to solve this?