def cost(self, y, y_true):
dijkstra = shortest_path.Dijkstra()
path = tf.cast(dijkstra(y), dtype=tf.float32)
y_true = tf.cast(y_true, dtype=tf.float32)
return tf.math.reduce_sum(tf.math.reduce_sum(y_true * path, axis=-1),
axis=-1)
def shortest_path_fn(x):
return tf.cast(shortest_path.Dijkstra()(x), dtype=x.dtype)